Function _split_handler

backends/mlx/ops.py:1628–1666

Handle split operation with uniform chunk size. Splits a tensor into chunks of a given size along a dimension. The last chunk may be smaller if the dimension does not divide evenly. PyTorch: split(x, split_size, dim=0) We pass [split_size] to the interpreter, which computes the ac

(P: MLXProgramBuilder, n: Node)

    target=[torch.ops.aten.split.Tensor, torch.ops.aten.split_copy.Tensor]
)
def _split_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    """Handle split operation with uniform chunk size.

    Splits a tensor into chunks of a given size along a dimension.
    The last chunk may be smaller if the dimension does not divide evenly.

    PyTorch: split(x, split_size, dim=0)

    We pass [split_size] to the interpreter, which computes the actual
    chunk sizes based on the tensor dimension.
    """
    args = P.args(n)
    require_args(args, 2, 3, "aten.split")
    require_kwargs(P.kwargs(n), set(), "aten.split")
    x = args[0]
    split_size = args[1]
    dim = args[2] if len(args) > 2 else 0

    axis = dim if dim is not None else 0
    if axis < 0:
        x_meta = n.args[0].meta.get("val")
        if x_meta is None:
            raise RuntimeError("split: missing tensor metadata for negative axis")
        axis += len(x_meta.shape)

    # Create output slots for multi-output operation
    output_slots = P.make_or_get_slots(n)

    # Emit SplitNode - interpreter computes actual chunk sizes from split_size
    P.emit(
        SplitNode(
            x=P.slot_to_tid(x),
            outs=[P.slot_to_tid(s) for s in output_slots],
            sizes=[P.to_int_or_vid(split_size)],
            axis=axis,
        )
    )

    return output_slots


@REGISTRY.register(target=[torch.ops.aten.repeat.default])