Function _arange_handler

backends/mlx/ops.py:2082–2116

Handle arange with just stop, or (start, stop) or (start, stop, step). Supports both static (literal int) and dynamic (Slot from item()) values.

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.aten.arange.default])
def _arange_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    """Handle arange with just stop, or (start, stop) or (start, stop, step).

    Supports both static (literal int) and dynamic (Slot from item()) values.
    """
    args = P.args(n)
    kwargs = P.kwargs(n)
    require_args(args, 1, 3, "aten.arange")
    require_kwargs(kwargs, {"dtype", "layout", "device", "pin_memory"}, "aten.arange")
    require_contiguous_format(
        layout=kwargs.get("layout"),
        op_name="aten.arange",
    )
    if len(args) == 1:
        start = 0
        stop = args[0]
    else:
        start, stop = args[0:2]
    step = args[2] if len(args) > 2 else 1

    # arange defaults to int64 when dtype is not specified (like torch.arange)
    dtype = kwargs.get("dtype", torch.int64)
    scalar_type_val = torch_dtype_to_scalar_type(dtype)

    out = P.make_or_get_slot(n)
    P.emit(
        ARangeNode(
            out=P.slot_to_tid(out),
            start=P.to_int_or_vid(start),
            stop=P.to_int_or_vid(stop),
            step=P.to_int_or_vid(step),
            scalar_type=scalar_type_val,
        )
    )
    return out


@REGISTRY.register(target=[torch.ops.aten.arange.start_step])