Function linspace

imperative/python/megengine/xla/rules/tensor.py:246–293

(
    start: Union[int, float, HLOTensor],
    stop: Union[int, float, HLOTensor],
    num: Union[int, HLOTensor],
    dtype=np.float32,
    oshape=None,
)

def linspace(
    start: Union[int, float, HLOTensor],
    stop: Union[int, float, HLOTensor],
    num: Union[int, HLOTensor],
    dtype=np.float32,
    oshape=None,
):
    if isinstance(num, HLOTensor):
        assert (
            oshape is not None
        ), "if num is a HLOTensor, please specify the output shape with oshape"

    if isinstance(num, int) and oshape is not None:
        assert _shape_equal(
            oshape, (num,)
        ), f"shape error: shape {oshape} .vs num {num}"

    oshape = (num,) if oshape is None else oshape

    if any([isinstance(x, HLOTensor) for x in [start, stop, num]]):
        start, stop, num = [
            x if isinstance(x, HLOTensor) else HLOTensor(float(x))
            for x in [start, stop, num]
        ]
        start, stop, num = [
            x if x.dtype == np.float32 else x.astype(np.float32)
            for x in [start, stop, num]
        ]
        # we want to implement: scale = where(num==1, 0.0, (stop-start)/(num-1.0))
        # to avoid divided by zero, but xla where is implemented by mul and add, the
        # result is still NAN, so we add a huge number to num to simulate.
        # we use 1e32 rather than np.inf because np.inf * zero will also return NAN
        num_refactor = where(num == 1.0, 1e32, num)
        scale = (stop - start) / (num_refactor - 1.0)
        scale = where(num == 1.0, 0.0, scale)
        offset = start
    else:
        assert all([isinstance(x, (int, float)) for x in [start, stop, num]])
        if num == 1:
            scale = 0
        else:
            scale = (stop - start) / (num - 1.0)
        offset = float(start)

    fsrc = iota(np.float32, oshape, -1)
    fout = fsrc * scale + offset

    return fout if dtype == np.float32 else fout.astype(dtype)


def arange(