Function linspace_lower

imperative/python/megengine/xla/rules/tensor.py:807–812

(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])

@register_lower_rule(mops.Linspace)
def linspace_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]):
    assert len(args) == 3 and len(ctx.vars_in) == 3 and len(ctx.vars_out) == 1
    assert len(ctx.vars_out[0].shape) == 1, "linspace/arange only support 1D"

    odtype, oshape = ctx.vars_out[0].dtype, ctx.vars_out[0].shape
    return linspace(args[0], args[1], args[2], odtype, oshape)