fix: use explicit dim_order kwarg for _clone_dim_order output spec in SpecPropPass

For _clone_dim_order and _to_dim_order_copy, output layout is determined by
the op's dim_order argument (e.g. from clone(memory_format=channels_last)),
not by the input tensor. SpecPropPass was propagating from the input's
FakeTensor, so contiguous input produced a contiguous output spec while
the kernel received dim_order=[0,2,3,1], causing runtime InvalidArgument.

- dim_order_utils: add _is_dim_order_op_with_explicit_arg and
  get_explicit_output_dim_order(node) to read dim_order from kwargs
- spec_prop_pass: set output spec dim_order from explicit kwargs when
  present; otherwise keep propagating from primary input (format-preserving)

Fixes test_op_clone_dim_order_propagation and related memory format tests.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: NefAI

exir/passes/dim_order_utils.py

@@ -18,6 +18,22 @@
     exir_ops = None  # type: ignore[assignment]
 
 
+def _is_dim_order_op_with_explicit_arg(op: object) -> bool:
+    """True if the op takes an explicit dim_order kwarg (e.g. _clone_dim_order, _to_dim_order_copy)."""
+    if exir_ops is None:
+        return False
+    return op in (
+        exir_ops.edge.dim_order_ops._clone_dim_order.default,
+        exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+    ) or (
+        hasattr(exir_ops.edge.dim_order_ops._clone_dim_order, "out")
+        and op == exir_ops.edge.dim_order_ops._clone_dim_order.out
+    ) or (
+        hasattr(exir_ops.edge.dim_order_ops._to_dim_order_copy, "out")
+        and op == exir_ops.edge.dim_order_ops._to_dim_order_copy.out
+    )
+
+
 def _format_preserving_ops() -> Set[object]:
     """Build set of format-preserving ops (aten and edge dialect)."""
     ops: Set[object] = {
@@ -61,6 +77,25 @@ def dim_order_from_fake_tensor(t: torch.Tensor) -> Optional[List[int]]:
         return None
 
 
+def get_explicit_output_dim_order(
+    node: "torch.fx.Node",
+) -> Optional[List[int]]:
+    """
+    If the node is a dim_order op (_clone_dim_order, _to_dim_order_copy) with
+    an explicit dim_order in kwargs, return it. Otherwise return None so the
+    caller can propagate from the primary input (format-preserving).
+    """
+    if not _is_dim_order_op_with_explicit_arg(node.target):
+        return None
+    dim_order_val = node.kwargs.get("dim_order") if node.kwargs else None
+    if dim_order_val is None:
+        return None
+    if isinstance(dim_order_val, (list, tuple)) and len(dim_order_val) > 0:
+        if all(isinstance(i, int) for i in dim_order_val):
+            return list(dim_order_val)
+    return None
+
+
 def should_propagate_dim_order(op: object) -> bool:
     """True if the op is format-preserving and we should propagate primary input dim_order to out."""
     return op in FORMAT_PRESERVING_OPS

exir/passes/spec_prop_pass.py

@@ -14,6 +14,7 @@
 from executorch.exir.pass_base import ExportPass, ProxyValue
 from executorch.exir.passes.dim_order_utils import (
     dim_order_from_fake_tensor,
+    get_explicit_output_dim_order,
     should_propagate_dim_order,
 )
 from executorch.exir.tensor import TensorSpec
@@ -85,33 +86,35 @@ def get_spec(x):
                         and should_propagate_dim_order(node.target)
                         and node.args
                     ):
-                        # Propagate primary input dim_order for format-preserving ops (Fix #16032).
-                        # Handles both clone.out (out= kwarg) and clone.default (single output).
-                        self_val = node.args[0].meta.get("val")
-                        if self_val is not None:
-                            src_dim_order = dim_order_from_fake_tensor(self_val)
-                            if "out" in node.kwargs:
-                                out_arg = node.kwargs["out"]
-                                assert isinstance(
-                                    out_arg, torch.fx.Node
-                                ), (
-                                    f"Expected clone.out 'out' to be fx.Node, got {type(out_arg)}"
+                        # Output dim_order: use explicit kwargs for dim_order ops
+                        # (_clone_dim_order, _to_dim_order_copy), else propagate from input (Fix #16032).
+                        if "out" in node.kwargs:
+                            out_arg = node.kwargs["out"]
+                            assert isinstance(
+                                out_arg, torch.fx.Node
+                            ), (
+                                f"Expected clone.out 'out' to be fx.Node, got {type(out_arg)}"
+                            )
+                            out_spec = out_arg.meta.get("spec")
+                        else:
+                            out_spec = node.meta.get("spec")
+                            if out_spec is None and meta_val is not None:
+                                node.meta["spec"] = pytree.tree_map(
+                                    make_spec, meta_val
                                 )
-                                out_spec = out_arg.meta.get("spec")
+                                out_spec = node.meta["spec"]
+                        if out_spec is not None and hasattr(out_spec, "dim_order"):
+                            explicit_dim_order = get_explicit_output_dim_order(node)
+                            if explicit_dim_order is not None:
+                                out_spec.dim_order = tuple(explicit_dim_order)
                             else:
-                                # clone.default: ensure node has spec (ExportPass may not set it)
-                                out_spec = node.meta.get("spec")
-                                if out_spec is None and meta_val is not None:
-                                    node.meta["spec"] = pytree.tree_map(
-                                        make_spec, meta_val
+                                self_val = node.args[0].meta.get("val")
+                                if self_val is not None:
+                                    src_dim_order = dim_order_from_fake_tensor(
+                                        self_val
                                     )
-                                    out_spec = node.meta["spec"]
-                            if (
-                                out_spec is not None
-                                and hasattr(out_spec, "dim_order")
-                                and src_dim_order is not None
-                            ):
-                                out_spec.dim_order = tuple(src_dim_order)
+                                    if src_dim_order is not None:
+                                        out_spec.dim_order = tuple(src_dim_order)
                     elif (
                         node.op == "call_function"
                         and node.target == executorch_call_delegate