Function iota

imperative/python/megengine/xla/rules/tensor.py:231–243

do some thing like arange. for example: shape = (2, 3), dimension=1, output is [[0, 1, 2], [0, 1, 2]] shape = (2, 3), dimension=-1, output is [[0, 0, 0], [1, 1, 1]]

(dtype, shape, dimension)

def iota(dtype, shape, dimension):
    """
    do some thing like arange.
    for example:
        shape = (2, 3), dimension=1, output is [[0, 1, 2], [0, 1, 2]]
        shape = (2, 3), dimension=-1, output is [[0, 0, 0], [1, 1, 1]]
    """
    dimension = dimension + len(shape) if dimension < 0 else dimension
    ret = hlo.IotaOp(
        ir_utils.make_ir_type_according_meta(shape, dtype), ir_utils.i64_attr(dimension)
    ).results
    assert len(ret) == 1, f"{len(ret)}"
    return HLOTensor(ret[0])


def linspace(