Function _arange_start_step_handler

backends/mlx/ops.py:2120–2153

Handle arange with start, end, and step arguments. Supports both static (literal int) and dynamic (Slot from item()) start/stop/step.

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.aten.arange.start_step])
def _arange_start_step_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    """Handle arange with start, end, and step arguments.

    Supports both static (literal int) and dynamic (Slot from item()) start/stop/step.
    """
    args = P.args(n)
    kwargs = P.kwargs(n)
    require_args(args, 2, 3, "aten.arange.start_step")
    require_kwargs(
        kwargs, {"dtype", "layout", "device", "pin_memory"}, "aten.arange.start_step"
    )
    require_contiguous_format(
        layout=kwargs.get("layout"),
        op_name="aten.arange.start_step",
    )
    start = args[0]
    stop = args[1]
    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.rms_norm.default])