Function _to_copy_handler

backends/mlx/ops.py:1037–1072

Handle aten._to_copy - lower-level dtype/device conversion.

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.aten._to_copy.default])
def _to_copy_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    """Handle aten._to_copy - lower-level dtype/device conversion."""
    # aten._to_copy(Tensor self, *, ScalarType? dtype=None, ...)
    args = P.args(n)
    kwargs = P.kwargs(n)
    require_args(args, 1, 1, "aten._to_copy")
    require_kwargs(
        kwargs, {"dtype", "device", "layout", "memory_format"}, "aten._to_copy"
    )
    require_contiguous_format(
        layout=kwargs.get("layout"),
        memory_format=kwargs.get("memory_format"),
        op_name="aten._to_copy",
    )
    x = args[0]
    out = P.make_or_get_slot(n)

    dtype = kwargs.get("dtype")
    if dtype is not None:
        # Dtype conversion
        P.emit(
            AsTypeNode(
                x=P.slot_to_tid(x),
                out=P.slot_to_tid(out),
                scalar_type=torch_dtype_to_scalar_type(dtype),
            )
        )
    else:
        # No dtype change, just copy (use contiguous)
        P.emit(
            ContiguousNode(
                x=P.slot_to_tid(x),
                out=P.slot_to_tid(out),
            )
        )
    return out


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