(P: MLXProgramBuilder, n: Node)
| 993 | # This is what x.to(dtype) becomes after to_edge() transformation |
| 994 | @REGISTRY.register(target=[exir_ops.edge.dim_order_ops._to_dim_order_copy.default]) |
| 995 | def _dim_order_copy_handler(P: MLXProgramBuilder, n: Node) -> Slot: |
| 996 | # dim_order_ops._to_dim_order_copy(Tensor self, *, ScalarType? dtype=None, ...) |
| 997 | # If dtype is specified, this is a dtype conversion (use AsTypeNode) |
| 998 | # If dtype is None/same, this is just a memory layout copy (use ContiguousNode) |
| 999 | args = P.args(n) |
| 1000 | kwargs = P.kwargs(n) |
| 1001 | require_args(args, 1, 1, "dim_order_ops._to_dim_order_copy") |
| 1002 | require_kwargs( |
| 1003 | kwargs, |
| 1004 | {"dtype", "device", "layout", "non_blocking", "dim_order"}, |
| 1005 | "dim_order_ops._to_dim_order_copy", |
| 1006 | ) |
| 1007 | require_contiguous_format( |
| 1008 | layout=kwargs.get("layout"), |
| 1009 | dim_order=kwargs.get("dim_order"), |
| 1010 | op_name="dim_order_ops._to_dim_order_copy", |
| 1011 | ) |
| 1012 | x = args[0] |
| 1013 | out = P.make_or_get_slot(n) |
| 1014 | |
| 1015 | dtype = kwargs.get("dtype") |
| 1016 | if dtype is not None: |
| 1017 | # Dtype conversion |
| 1018 | P.emit( |
| 1019 | AsTypeNode( |
| 1020 | x=P.slot_to_tid(x), |
| 1021 | out=P.slot_to_tid(out), |
| 1022 | scalar_type=torch_dtype_to_scalar_type(dtype), |
| 1023 | ) |
| 1024 | ) |
| 1025 | else: |
| 1026 | # No dtype change, just memory layout (contiguous) |
| 1027 | P.emit( |
| 1028 | ContiguousNode( |
| 1029 | x=P.slot_to_tid(x), |
| 1030 | out=P.slot_to_tid(out), |
| 1031 | ) |
| 1032 | ) |
| 1033 | return out |
| 1034 | |
| 1035 | |
| 1036 | @REGISTRY.register(target=[torch.ops.aten._to_copy.default]) |
nothing calls this directly
no test coverage detected