Function _slice_scatter_handler

backends/mlx/ops.py:1863–1900

Handle aten.slice_scatter: return a copy of self with self[dim][start:end:step] = src.

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.aten.slice_scatter.default])
def _slice_scatter_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    """Handle aten.slice_scatter: return a copy of self with self[dim][start:end:step] = src."""
    args = P.args(n)
    require_args(args, 2, 6, "aten.slice_scatter")
    require_kwargs(P.kwargs(n), set(), "aten.slice_scatter")
    self_tensor = args[0]
    src = args[1]
    dim = args[2] if len(args) > 2 else 0
    start = args[3] if len(args) > 3 else 0
    end = args[4] if len(args) > 4 else None
    step = args[5] if len(args) > 5 else 1

    # If end is None, default to dim size
    if end is None:
        input_meta = n.args[0].meta.get("val")
        if input_meta is not None:
            end = input_meta.shape[dim]
        else:
            raise ValueError(
                "aten.slice_scatter: end=None requires input shape metadata"
            )

    require_static_int(step, "step", "aten.slice_scatter")
    assert step >= 1, f"aten.slice_scatter: step must be >= 1, got {step}"

    out = P.make_or_get_slot(n)
    P.emit(
        SliceUpdateNode(
            dst=P.slot_to_tid(self_tensor),
            update=P.slot_to_tid(src),
            out=P.slot_to_tid(out),
            axis=P.to_int_or_vid(dim),
            start=P.to_int_or_vid(start),
            stop=P.to_int_or_vid(end),
            step=step,
        )
    )
    return out


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