From d662f2fed5f129f83473f6fd72cb54c37376ab5f Mon Sep 17 00:00:00 2001
From: Jake Stevens <jrstevens@meta.com>
Date: Fri, 11 Sep 2026 07:09:25 -0700
Subject: [PATCH 1/2] Fix dim_order derivation for channels-last tensors with
 size-1 dims

A channels-last tensor with a size-1 channel (e.g. the (1, 1, 49, 10)
input of MLPerf Tiny keyword spotting) has strides whose stable sort is
the non-canonical dim order (0, 2, 1, 3). dim_order_from_stride now takes
the tensor sizes and returns the canonical dim order when the strides
exactly match a contiguous or channels-last layout, which portable
kernels require. The canonical order describes the identical physical
layout, and genuinely non-canonical layouts still fall back to sorting.

Fixes https://github.com/pytorch/executorch/issues/22520
---
 .../native/serialization/graph_serialize.py   |  4 +-
 backends/native/test/test_serialize.py        |  6 +-
 exir/emit/_emitter.py                         |  2 +-
 exir/passes/memory_format_ops_pass.py         | 12 ++-
 exir/tensor.py                                | 88 ++++++++++++++++++-
 exir/tensor_layout.py                         |  2 +-
 exir/tests/test_dim_order_utils.py            | 48 +++++++++-
 exir/tests/test_tensor.py                     | 44 ++++++++++
 8 files changed, 196 insertions(+), 10 deletions(-)

diff --git a/backends/native/serialization/graph_serialize.py b/backends/native/serialization/graph_serialize.py
index 0d89842cabd..feb146bfae4 100644
--- a/backends/native/serialization/graph_serialize.py
+++ b/backends/native/serialization/graph_serialize.py
@@ -203,7 +203,9 @@ def _dim_order(t: torch.Tensor) -> list[int]:
     strides = tuple(t.stride())
     sizes = list(t.shape)
     # dim_order_from_stride handles symbolic strides and rejects stride-0 layouts.
-    dim_order = [int(d) for d in dim_order_from_stride(strides)]
+    # Sizes are passed so size-ambiguous canonical layouts (e.g. channels-last
+    # with a size-1 channel) serialize the canonical dim order.
+    dim_order = [int(d) for d in dim_order_from_stride(strides, tuple(sizes))]
     expected = stride_from_dim_order(sizes, dim_order)
     for i in range(ndim):
         # A size-1 dim only ever indexes 0, so its stride is arbitrary and need not
diff --git a/backends/native/test/test_serialize.py b/backends/native/test/test_serialize.py
index edb1c538c89..b31d03abdb7 100644
--- a/backends/native/test/test_serialize.py
+++ b/backends/native/test/test_serialize.py
@@ -555,9 +555,11 @@ def test_permuted_contiguous(self):
 
     def test_channels_last_with_size_one_channel(self):
         # channels-last leaves the size-1 channel with an arbitrary stride, which
-        # must not be treated as a non-expressible layout.
+        # must not be treated as a non-expressible layout. The size-ambiguous
+        # strides serialize to the canonical channels-last dim order, which
+        # describes the identical physical layout.
         t = torch.randn(2, 1, 3, 4).to(memory_format=torch.channels_last)
-        self.assertEqual(_dim_order(t), [0, 2, 1, 3])
+        self.assertEqual(_dim_order(t), [0, 2, 3, 1])
 
     def test_sliced_layout_raises(self):
         t = torch.randn(4, 8)[:, :4]
diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py
index 4e3baf26592..e78b860ac3b 100644
--- a/exir/emit/_emitter.py
+++ b/exir/emit/_emitter.py
@@ -2034,7 +2034,7 @@ def _is_buffer(node: Node, graph_signature: ExportGraphSignature) -> bool:
 
                 spec.storage = real_tensor.untyped_storage()
                 spec.stride = real_tensor.stride()
-                spec.dim_order = dim_order_from_stride(spec.stride)
+                spec.dim_order = dim_order_from_stride(spec.stride, tuple(spec.shape))
             # User inputs and mutable buffers are not constants, other buffers or parameters are.
             if initialize_buffer and is_mutable_buffer:
                 spec.const = True
diff --git a/exir/passes/memory_format_ops_pass.py b/exir/passes/memory_format_ops_pass.py
index 13468dfd8d8..b4e82b46fbd 100644
--- a/exir/passes/memory_format_ops_pass.py
+++ b/exir/passes/memory_format_ops_pass.py
@@ -16,6 +16,7 @@
     DimOrderOpsMap,
     MemoryFormatOpsMap,
 )
+from executorch.exir.tensor import dim_order_from_stride
 
 logger = logging.getLogger(__file__)
 logger.setLevel(logging.INFO)
@@ -64,9 +65,16 @@ def call_operator(self, op, args, kwargs, meta):
         # Derive dim_order based on memory format
         dim_order: List[int]
         if mem_format in (None, torch.preserve_format):
-            # preserve_format: inherit dim_order from input tensor
+            # preserve_format: inherit dim_order from input tensor,
+            # canonicalized so size-1 dims keep the canonical
+            # contiguous/channels-last dim order.
             if input_tensor is not None:
-                dim_order = [int(d) for d in input_tensor.dim_order()]
+                dim_order = [
+                    int(d)
+                    for d in dim_order_from_stride(
+                        input_tensor.stride(), tuple(input_tensor.shape)
+                    )
+                ]
             else:
                 # Fallback to contiguous if no single input tensor is available
                 # (e.g. list inputs like torch.stack).
diff --git a/exir/tensor.py b/exir/tensor.py
index 199c5adfafe..7d75dda02a9 100644
--- a/exir/tensor.py
+++ b/exir/tensor.py
@@ -49,7 +49,79 @@ def contiguous_stride_from_shape(shape: torch.Size) -> Tuple[int]:
     return tuple(reversed(strides))
 
 
-def dim_order_from_stride(stride: Tuple[int]) -> Tuple[bytes]:
+def _expected_channels_last_stride(sizes: Tuple[int]) -> Optional[Tuple[int]]:
+    """
+    Expected strides of a channels-last tensor with the given sizes.
+
+    Only 4D (NCHW) and 5D (NCDHW) tensors have a channels-last layout;
+    returns None for any other rank.
+    """
+    ndim = len(sizes)
+    if ndim not in (4, 5):
+        return None
+    # Stride of dim 1 (channels) is 1, stride of dim i >= 2 is
+    # channels * prod(sizes[i + 1 :]), stride of dim 0 is prod(sizes[1:]).
+    expected = [1] * ndim
+    accum = sizes[1]
+    for i in range(ndim - 1, 1, -1):
+        expected[i] = accum
+        accum = accum * sizes[i]
+    expected[0] = accum
+    return tuple(expected)
+
+
+def _strides_match(stride: Tuple[int], expected: Tuple[int]) -> bool:
+    """
+    Exact stride equality that is safe for symbolic (SymInt) values.
+
+    Uses statically_known_true so no shape guards are added; returns False
+    when equality cannot be proven (or comparison itself fails), in which
+    case callers fall back to sorting.
+    """
+    if len(stride) != len(expected):
+        return False
+    pairs = list(zip(stride, expected))
+    if all(type(s) is int and type(e) is int for s, e in pairs):
+        return all(s == e for s, e in pairs)
+    from torch.fx.experimental.symbolic_shapes import statically_known_true
+
+    try:
+        return all(statically_known_true(s == e) for s, e in pairs)
+    except Exception:
+        return False
+
+
+def _canonical_dim_order_for_strides(
+    stride: Tuple[int], sizes: Tuple[int]
+) -> Optional[Tuple[bytes]]:
+    """
+    Canonical dim order when the strides exactly match a canonical layout.
+
+    Returns None when the strides match neither contiguous nor channels-last
+    (or when sizes/strides are symbolic and equality cannot be proven), in
+    which case the caller falls back to sorting strides.
+
+    Contiguous is checked first so that shapes whose strides match both
+    layouts (only possible when every dim but the batch is size 1, e.g.
+    (N, 1, 1, 1)) keep the historical stable-sort result, matching
+    torch.Tensor.dim_order().
+    """
+    ndim = len(stride)
+    if ndim == 0:
+        return ()
+    if _strides_match(stride, contiguous_stride_from_shape(torch.Size(sizes))):
+        return tuple(typing.cast(Tuple[bytes], tuple(range(ndim))))
+    if ndim in (4, 5):
+        expected = _expected_channels_last_stride(sizes)
+        assert expected is not None
+        if _strides_match(stride, expected):
+            return tuple(typing.cast(Tuple[bytes], (0, *range(2, ndim), 1)))
+    return None
+
+
+def dim_order_from_stride(
+    stride: Tuple[int], sizes: Optional[Tuple[int]] = None
+) -> Tuple[bytes]:
     """
     Dimension order represents how dimensions are laid out in memory,
     starting from the outer-most to the inner-most dimension.
@@ -65,12 +137,24 @@ def dim_order_from_stride(stride: Tuple[int]) -> Tuple[bytes]:
     in original order. Thus when strides = (4, 3, 1, 1) returned value is (0, 1, 2, 3)
     Another example is: sizes = (1, 3, 1, 1) with strides = (3, 1, 3, 3), returned
     value is (0, 2, 3, 1)
+
+    When sizes are provided and the strides exactly match the strides of a
+    contiguous or channels-last tensor of that shape, the canonical dim order
+    is returned instead. Size-1 dimensions make the sort ambiguous: e.g. a
+    channels-last (N, 1, H, W) tensor has strides (H*W, 1, W, 1), whose stable
+    sort is the non-canonical (0, 2, 1, 3) that portable kernels reject, even
+    though the canonical (0, 2, 3, 1) describes the identical physical layout.
     """
     from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true
 
     for s in stride:
         torch._check(s != 0, lambda: "0 in strides is not supported for ExecuTorch.")
 
+    if sizes is not None and len(sizes) == len(stride):
+        canonical = _canonical_dim_order_for_strides(stride, sizes)
+        if canonical is not None:
+            return canonical
+
     class K(NamedTuple):
         stride: int
 
@@ -201,7 +285,7 @@ def from_tensor(cls, tensor: torch.Tensor, const: bool = False) -> TensorSpec:
             is_sparse=tensor.is_sparse,
         )
         spec.stride = tensor.stride()
-        spec.dim_order = dim_order_from_stride(spec.stride)
+        spec.dim_order = dim_order_from_stride(spec.stride, tuple(spec.shape))
         spec.requires_grad = tensor.requires_grad
         spec.storage = tensor.untyped_storage() if const else None
 
diff --git a/exir/tensor_layout.py b/exir/tensor_layout.py
index e902006783f..dfcdc9b3a74 100644
--- a/exir/tensor_layout.py
+++ b/exir/tensor_layout.py
@@ -35,5 +35,5 @@ def from_tensor(cls, tensor: torch.Tensor) -> "TensorLayout":
         return TensorLayout(
             scalar_type=scalar_type_enum(tensor.dtype),
             sizes=list(tensor.shape),
-            dim_order=list(dim_order_from_stride(tensor.stride())),
+            dim_order=list(dim_order_from_stride(tensor.stride(), tuple(tensor.shape))),
         )
diff --git a/exir/tests/test_dim_order_utils.py b/exir/tests/test_dim_order_utils.py
index 6a8ea934db0..0926ad4fc40 100644
--- a/exir/tests/test_dim_order_utils.py
+++ b/exir/tests/test_dim_order_utils.py
@@ -7,9 +7,12 @@
 # pyre-strict
 import unittest
 
+from typing import List
+
 import torch
-from executorch.exir import to_edge_transform_and_lower
+from executorch.exir import to_edge, to_edge_transform_and_lower
 from executorch.exir.dim_order_utils import get_dim_order, get_memory_format
+from executorch.exir.schema import KernelCall, Tensor
 
 
 class TestDimOrderUtils(unittest.TestCase):
@@ -47,3 +50,46 @@ def forward(self, t1, t2):
         expo_prog = torch.export.export(M, (x, y))
         edge_prog = to_edge_transform_and_lower(expo_prog)
         edge_prog.to_executorch()
+
+    def test_channels_last_single_channel_conv_dim_order(self) -> None:
+        # Regression test for https://github.com/pytorch/executorch/issues/22520:
+        # a channels-last conv input with a size-1 channel must be emitted with
+        # the canonical channels-last dim order. Portable conv kernels reject
+        # anything but default/channels-last dim orders, and require the input
+        # and output dim orders to match.
+        class Conv(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = torch.nn.Conv2d(1, 8, kernel_size=3, padding=1)
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return self.conv(x)
+
+        model = Conv().to(memory_format=torch.channels_last).eval()
+        inputs = (torch.randn(1, 1, 12, 12).to(memory_format=torch.channels_last),)
+        edge_prog = to_edge(torch.export.export(model, inputs, strict=True))
+        et_program = edge_prog.to_executorch().executorch_program
+
+        conv_found = False
+        for inst in et_program.execution_plan[0].chains[0].instructions:
+            kernel = inst.instr_args
+            if not isinstance(kernel, KernelCall):
+                continue
+            op = et_program.execution_plan[0].operators[kernel.op_index]
+            if op.name != "aten::convolution":
+                continue
+            conv_found = True
+            tensors: List[Tensor] = []
+            for arg in dict.fromkeys(kernel.args):
+                val = et_program.execution_plan[0].values[arg].val
+                if isinstance(val, Tensor) and len(val.sizes) == 4:
+                    tensors.append(val)
+            # Input, weight, and output.
+            self.assertEqual(len(tensors), 3)
+            for tensor in tensors:
+                self.assertIn(
+                    list(tensor.dim_order),
+                    ([0, 1, 2, 3], [0, 2, 3, 1]),
+                )
+            self.assertEqual(list(tensors[0].dim_order), list(tensors[-1].dim_order))
+        self.assertTrue(conv_found)
diff --git a/exir/tests/test_tensor.py b/exir/tests/test_tensor.py
index 19a536d44d5..3e7e979744e 100644
--- a/exir/tests/test_tensor.py
+++ b/exir/tests/test_tensor.py
@@ -321,6 +321,50 @@ def test_dim_order_from_stride_unbacked(self) -> None:
         with self.assertRaises(RuntimeError):
             dim_order_from_stride((u0, 0, 1))
 
+    def test_dim_order_from_stride_with_sizes(self) -> None:
+        # Regression test for https://github.com/pytorch/executorch/issues/22520:
+        # a channels-last tensor with a size-1 channel has strides whose stable
+        # sort is the non-canonical (0, 2, 1, 3), which portable kernels reject.
+        # With sizes, the canonical dim order is returned instead.
+        dim_order = dim_order_from_stride((490, 1, 10, 1), (1, 1, 49, 10))
+        self.assertEqual((0, 2, 3, 1), dim_order)
+        # The canonical dim order describes the identical physical layout.
+        self.assertEqual(
+            [490, 1, 10, 1],
+            stride_from_dim_order([1, 1, 49, 10], [0, 2, 3, 1]),
+        )
+
+        # Contiguous tensors are unaffected, including with size-1 dims.
+        self.assertEqual(
+            (0, 1, 2, 3),
+            dim_order_from_stride((490, 490, 10, 1), (1, 1, 49, 10)),
+        )
+
+        # Shapes matching both canonical layouts keep the historical result.
+        self.assertEqual(
+            (0, 1, 2, 3), dim_order_from_stride((1, 1, 1, 1), (2, 1, 1, 1))
+        )
+
+        # 5D channels-last with a size-1 channel.
+        self.assertEqual(
+            (0, 2, 3, 4, 1),
+            dim_order_from_stride((120, 1, 30, 6, 1), (1, 1, 4, 5, 6)),
+        )
+
+        # Genuinely non-canonical layouts still fall back to sorting.
+        self.assertEqual(
+            (3, 1, 2, 0),
+            dim_order_from_stride((1, 20, 5, 60), (2, 3, 4, 5)),
+        )
+
+        # Without sizes, behavior is unchanged (stable sort).
+        self.assertEqual((0, 2, 1, 3), dim_order_from_stride((490, 1, 10, 1)))
+
+        # TensorSpec picks up the canonical dim order (used by SpecPropPass).
+        t = torch.empty(1, 1, 49, 10).to(memory_format=torch.channels_last)
+        spec = TensorSpec.from_tensor(t)
+        self.assertEqual((0, 2, 3, 1), spec.dim_order)
+
     def test_strides_from_dim_order(self) -> None:
         sizes = []
         dim_order = []

From 66bd993b4d9fb9511df7efaaf6e078ce0218b986 Mon Sep 17 00:00:00 2001
From: Jake Stevens <jrstevens@meta.com>
Date: Fri, 11 Sep 2026 09:11:16 -0700
Subject: [PATCH 2/2] Preserve supported dim orders when canonicalizing strides

Reuse PyTorch stride calculation while keeping existing default orders and preserve_format inference. Cover singleton-layout regressions and symbolic sizes, and restore portable NXP KWS comparisons with finite tolerances.

Authored with AI assistance (Codex).
---
 .../native/serialization/graph_serialize.py   |   2 -
 backends/native/test/test_serialize.py        |   5 +-
 .../test_mlperf_tiny_keyword_spotting.py      |  30 +-----
 exir/passes/memory_format_ops_pass.py         |  12 +--
 exir/tensor.py                                | 102 ++++--------------
 exir/tests/test_dim_order_utils.py            |  41 ++++++-
 exir/tests/test_tensor.py                     |  41 +++++--
 7 files changed, 93 insertions(+), 140 deletions(-)

diff --git a/backends/native/serialization/graph_serialize.py b/backends/native/serialization/graph_serialize.py
index feb146bfae4..1f9a333c871 100644
--- a/backends/native/serialization/graph_serialize.py
+++ b/backends/native/serialization/graph_serialize.py
@@ -203,8 +203,6 @@ def _dim_order(t: torch.Tensor) -> list[int]:
     strides = tuple(t.stride())
     sizes = list(t.shape)
     # dim_order_from_stride handles symbolic strides and rejects stride-0 layouts.
-    # Sizes are passed so size-ambiguous canonical layouts (e.g. channels-last
-    # with a size-1 channel) serialize the canonical dim order.
     dim_order = [int(d) for d in dim_order_from_stride(strides, tuple(sizes))]
     expected = stride_from_dim_order(sizes, dim_order)
     for i in range(ndim):
diff --git a/backends/native/test/test_serialize.py b/backends/native/test/test_serialize.py
index b31d03abdb7..0d833968567 100644
--- a/backends/native/test/test_serialize.py
+++ b/backends/native/test/test_serialize.py
@@ -554,10 +554,7 @@ def test_permuted_contiguous(self):
         self.assertEqual(_dim_order(t), [0, 2, 1])
 
     def test_channels_last_with_size_one_channel(self):
-        # channels-last leaves the size-1 channel with an arbitrary stride, which
-        # must not be treated as a non-expressible layout. The size-ambiguous
-        # strides serialize to the canonical channels-last dim order, which
-        # describes the identical physical layout.
+        # A size-1 channel must not make channels-last non-expressible.
         t = torch.randn(2, 1, 3, 4).to(memory_format=torch.channels_last)
         self.assertEqual(_dim_order(t), [0, 2, 3, 1])
 
diff --git a/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py b/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py
index ab56849e970..22fb84d66be 100644
--- a/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py
+++ b/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py
@@ -29,11 +29,11 @@
 
 BOUNDS_MSE = {
     "PTQ": {
-        "channels-last": np.inf,
+        "channels-last": 3.5e-4,
         "channels-first": 5.5e-7,
     },
     "QAT": {
-        "channels-last": np.inf,
+        "channels-last": 3.5e-4,
         "channels-first": 3.3e-5,
     },
 }
@@ -45,19 +45,7 @@ def reseed_model_per_test_run():
     np.random.seed(23)
 
 
-@pytest.mark.parametrize(
-    "channels_last",
-    [
-        False,
-        pytest.param(
-            True,
-            marks=pytest.mark.xfail(
-                reason="EIEX-1082, don't forget to readjust bounds when it start working",
-                strict=True,
-            ),
-        ),
-    ],
-)
+@pytest.mark.parametrize("channels_last", [False, True])
 def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last, use_qat):
     # approx. 5 samples per class
     num_samples = 60
@@ -87,16 +75,6 @@ def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last, use_qat)
         partial(kws.train_model_fn, channels_last=channels_last) if use_qat else None
     )
 
-    # This model does not work in channels-last format when running with portable kernels.
-    # See more information below.
-    # Github issue: https://github.com/pytorch/executorch/issues/22520
-    # NXP internal issue ID: EIEX-1074
-    ref_model = (
-        ReferenceModel.QUANTIZED_EDGE_PYTHON
-        if channels_last
-        else ReferenceModel.QUANTIZED_EXECUTORCH_CPP
-    )
-
     lower_run_compare(
         model,
         [input_spec],
@@ -105,7 +83,7 @@ def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last, use_qat)
         dataset_creator=dataset_creator,
         output_comparator=comparator,
         mocker=mocker,
-        reference_model=ref_model,
+        reference_model=ReferenceModel.QUANTIZED_EXECUTORCH_CPP,
         use_qat=use_qat,
         train_fn=train_fn,
     )
diff --git a/exir/passes/memory_format_ops_pass.py b/exir/passes/memory_format_ops_pass.py
index b4e82b46fbd..13468dfd8d8 100644
--- a/exir/passes/memory_format_ops_pass.py
+++ b/exir/passes/memory_format_ops_pass.py
@@ -16,7 +16,6 @@
     DimOrderOpsMap,
     MemoryFormatOpsMap,
 )
-from executorch.exir.tensor import dim_order_from_stride
 
 logger = logging.getLogger(__file__)
 logger.setLevel(logging.INFO)
@@ -65,16 +64,9 @@ def call_operator(self, op, args, kwargs, meta):
         # Derive dim_order based on memory format
         dim_order: List[int]
         if mem_format in (None, torch.preserve_format):
-            # preserve_format: inherit dim_order from input tensor,
-            # canonicalized so size-1 dims keep the canonical
-            # contiguous/channels-last dim order.
+            # preserve_format: inherit dim_order from input tensor
             if input_tensor is not None:
-                dim_order = [
-                    int(d)
-                    for d in dim_order_from_stride(
-                        input_tensor.stride(), tuple(input_tensor.shape)
-                    )
-                ]
+                dim_order = [int(d) for d in input_tensor.dim_order()]
             else:
                 # Fallback to contiguous if no single input tensor is available
                 # (e.g. list inputs like torch.stack).
diff --git a/exir/tensor.py b/exir/tensor.py
index 7d75dda02a9..50fcd46fa53 100644
--- a/exir/tensor.py
+++ b/exir/tensor.py
@@ -49,76 +49,6 @@ def contiguous_stride_from_shape(shape: torch.Size) -> Tuple[int]:
     return tuple(reversed(strides))
 
 
-def _expected_channels_last_stride(sizes: Tuple[int]) -> Optional[Tuple[int]]:
-    """
-    Expected strides of a channels-last tensor with the given sizes.
-
-    Only 4D (NCHW) and 5D (NCDHW) tensors have a channels-last layout;
-    returns None for any other rank.
-    """
-    ndim = len(sizes)
-    if ndim not in (4, 5):
-        return None
-    # Stride of dim 1 (channels) is 1, stride of dim i >= 2 is
-    # channels * prod(sizes[i + 1 :]), stride of dim 0 is prod(sizes[1:]).
-    expected = [1] * ndim
-    accum = sizes[1]
-    for i in range(ndim - 1, 1, -1):
-        expected[i] = accum
-        accum = accum * sizes[i]
-    expected[0] = accum
-    return tuple(expected)
-
-
-def _strides_match(stride: Tuple[int], expected: Tuple[int]) -> bool:
-    """
-    Exact stride equality that is safe for symbolic (SymInt) values.
-
-    Uses statically_known_true so no shape guards are added; returns False
-    when equality cannot be proven (or comparison itself fails), in which
-    case callers fall back to sorting.
-    """
-    if len(stride) != len(expected):
-        return False
-    pairs = list(zip(stride, expected))
-    if all(type(s) is int and type(e) is int for s, e in pairs):
-        return all(s == e for s, e in pairs)
-    from torch.fx.experimental.symbolic_shapes import statically_known_true
-
-    try:
-        return all(statically_known_true(s == e) for s, e in pairs)
-    except Exception:
-        return False
-
-
-def _canonical_dim_order_for_strides(
-    stride: Tuple[int], sizes: Tuple[int]
-) -> Optional[Tuple[bytes]]:
-    """
-    Canonical dim order when the strides exactly match a canonical layout.
-
-    Returns None when the strides match neither contiguous nor channels-last
-    (or when sizes/strides are symbolic and equality cannot be proven), in
-    which case the caller falls back to sorting strides.
-
-    Contiguous is checked first so that shapes whose strides match both
-    layouts (only possible when every dim but the batch is size 1, e.g.
-    (N, 1, 1, 1)) keep the historical stable-sort result, matching
-    torch.Tensor.dim_order().
-    """
-    ndim = len(stride)
-    if ndim == 0:
-        return ()
-    if _strides_match(stride, contiguous_stride_from_shape(torch.Size(sizes))):
-        return tuple(typing.cast(Tuple[bytes], tuple(range(ndim))))
-    if ndim in (4, 5):
-        expected = _expected_channels_last_stride(sizes)
-        assert expected is not None
-        if _strides_match(stride, expected):
-            return tuple(typing.cast(Tuple[bytes], (0, *range(2, ndim), 1)))
-    return None
-
-
 def dim_order_from_stride(
     stride: Tuple[int], sizes: Optional[Tuple[int]] = None
 ) -> Tuple[bytes]:
@@ -138,23 +68,19 @@ def dim_order_from_stride(
     Another example is: sizes = (1, 3, 1, 1) with strides = (3, 1, 3, 3), returned
     value is (0, 2, 3, 1)
 
-    When sizes are provided and the strides exactly match the strides of a
-    contiguous or channels-last tensor of that shape, the canonical dim order
-    is returned instead. Size-1 dimensions make the sort ambiguous: e.g. a
-    channels-last (N, 1, H, W) tensor has strides (H*W, 1, W, 1), whose stable
-    sort is the non-canonical (0, 2, 1, 3) that portable kernels reject, even
-    though the canonical (0, 2, 3, 1) describes the identical physical layout.
+    With sizes, ambiguous non-canonical orders are corrected when strides
+    exactly match channels-last. Preserve existing default orders: portable
+    kernels require matching dim orders even for physically equivalent layouts.
     """
-    from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true
+    from torch.fx.experimental.symbolic_shapes import (
+        guard_or_false,
+        guard_or_true,
+        statically_known_true,
+    )
 
     for s in stride:
         torch._check(s != 0, lambda: "0 in strides is not supported for ExecuTorch.")
 
-    if sizes is not None and len(sizes) == len(stride):
-        canonical = _canonical_dim_order_for_strides(stride, sizes)
-        if canonical is not None:
-            return canonical
-
     class K(NamedTuple):
         stride: int
 
@@ -177,6 +103,18 @@ def __lt__(self, other):
     sorted_dims = [
         i[0] for i in sorted(enumerate(stride), key=lambda x: K(x[1]), reverse=True)
     ]
+    ndim = len(stride)
+    if (
+        sizes is not None
+        and len(sizes) == ndim
+        and ndim in (4, 5)
+        and sorted_dims != list(range(ndim))
+    ):
+        from torch._prims_common import make_channels_last_strides_for
+
+        expected = make_channels_last_strides_for(sizes)
+        if all(statically_known_true(s == e) for s, e in zip(stride, expected)):
+            sorted_dims = [0, *range(2, ndim), 1]
     return tuple(typing.cast(Tuple[bytes], sorted_dims))
 
 
diff --git a/exir/tests/test_dim_order_utils.py b/exir/tests/test_dim_order_utils.py
index 0926ad4fc40..44a65025e19 100644
--- a/exir/tests/test_dim_order_utils.py
+++ b/exir/tests/test_dim_order_utils.py
@@ -52,11 +52,6 @@ def forward(self, t1, t2):
         edge_prog.to_executorch()
 
     def test_channels_last_single_channel_conv_dim_order(self) -> None:
-        # Regression test for https://github.com/pytorch/executorch/issues/22520:
-        # a channels-last conv input with a size-1 channel must be emitted with
-        # the canonical channels-last dim order. Portable conv kernels reject
-        # anything but default/channels-last dim orders, and require the input
-        # and output dim orders to match.
         class Conv(torch.nn.Module):
             def __init__(self) -> None:
                 super().__init__()
@@ -93,3 +88,39 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                 )
             self.assertEqual(list(tensors[0].dim_order), list(tensors[-1].dim_order))
         self.assertTrue(conv_found)
+
+    def test_singleton_dims_preserve_default_order(self) -> None:
+        class Add(torch.nn.Module):
+            def forward(self, x, y):
+                return x + y
+
+        cases = (
+            (
+                Add(),
+                (
+                    torch.randn(2, 1, 3, 1).to(memory_format=torch.channels_last),
+                    torch.randn(2, 1, 3, 1),
+                ),
+            ),
+            (
+                torch.nn.ReLU(),
+                (torch.randn(2, 1, 3, 1, 1).to(memory_format=torch.channels_last_3d),),
+            ),
+        )
+        for model, inputs in cases:
+            with self.subTest(model=type(model).__name__):
+                program = (
+                    to_edge_transform_and_lower(torch.export.export(model, inputs))
+                    .to_executorch()
+                    .executorch_program
+                )
+                tensors = [
+                    value.val
+                    for value in program.execution_plan[0].values
+                    if isinstance(value.val, Tensor)
+                ]
+                self.assertEqual(len(tensors), len(inputs) + 1)
+                for tensor in tensors:
+                    self.assertEqual(
+                        list(tensor.dim_order), list(range(len(tensor.sizes)))
+                    )
diff --git a/exir/tests/test_tensor.py b/exir/tests/test_tensor.py
index 3e7e979744e..d6307d716e1 100644
--- a/exir/tests/test_tensor.py
+++ b/exir/tests/test_tensor.py
@@ -322,49 +322,68 @@ def test_dim_order_from_stride_unbacked(self) -> None:
             dim_order_from_stride((u0, 0, 1))
 
     def test_dim_order_from_stride_with_sizes(self) -> None:
-        # Regression test for https://github.com/pytorch/executorch/issues/22520:
-        # a channels-last tensor with a size-1 channel has strides whose stable
-        # sort is the non-canonical (0, 2, 1, 3), which portable kernels reject.
-        # With sizes, the canonical dim order is returned instead.
         dim_order = dim_order_from_stride((490, 1, 10, 1), (1, 1, 49, 10))
         self.assertEqual((0, 2, 3, 1), dim_order)
-        # The canonical dim order describes the identical physical layout.
         self.assertEqual(
             [490, 1, 10, 1],
             stride_from_dim_order([1, 1, 49, 10], [0, 2, 3, 1]),
         )
 
-        # Contiguous tensors are unaffected, including with size-1 dims.
         self.assertEqual(
             (0, 1, 2, 3),
             dim_order_from_stride((490, 490, 10, 1), (1, 1, 49, 10)),
         )
 
-        # Shapes matching both canonical layouts keep the historical result.
         self.assertEqual(
             (0, 1, 2, 3), dim_order_from_stride((1, 1, 1, 1), (2, 1, 1, 1))
         )
 
-        # 5D channels-last with a size-1 channel.
         self.assertEqual(
             (0, 2, 3, 4, 1),
             dim_order_from_stride((120, 1, 30, 6, 1), (1, 1, 4, 5, 6)),
         )
 
-        # Genuinely non-canonical layouts still fall back to sorting.
         self.assertEqual(
             (3, 1, 2, 0),
             dim_order_from_stride((1, 20, 5, 60), (2, 3, 4, 5)),
         )
 
-        # Without sizes, behavior is unchanged (stable sort).
         self.assertEqual((0, 2, 1, 3), dim_order_from_stride((490, 1, 10, 1)))
 
-        # TensorSpec picks up the canonical dim order (used by SpecPropPass).
         t = torch.empty(1, 1, 49, 10).to(memory_format=torch.channels_last)
         spec = TensorSpec.from_tensor(t)
         self.assertEqual((0, 2, 3, 1), spec.dim_order)
 
+    def test_dim_order_from_stride_preserves_supported_orders(self) -> None:
+        for sizes, strides in (
+            ((2, 1, 3, 1), (3, 1, 1, 1)),
+            ((2, 1, 3, 1, 1), (3, 1, 1, 1, 1)),
+            ((2, 3, 4, 5), (60, 1, 15, 3)),
+            ((2, 3, 4, 5, 6), (360, 1, 90, 18, 3)),
+        ):
+            with self.subTest(sizes=sizes):
+                self.assertEqual(
+                    dim_order_from_stride(strides),
+                    dim_order_from_stride(strides, sizes),
+                )
+
+    def test_dim_order_from_stride_with_symbolic_sizes(self) -> None:
+        from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+        shape_env = ShapeEnv()
+        height = shape_env.create_unbacked_symint()
+        torch._check_is_size(height)
+        torch._check(height >= 2)
+        self.assertEqual(
+            (0, 2, 3, 1),
+            dim_order_from_stride((10 * height, 1, 10, 1), (2, 1, height, 10)),
+        )
+        self.assertEqual(
+            (0, 2, 1, 3),
+            dim_order_from_stride((490, 1, 10, 1), (2, 1, height, 10)),
+        )
+        self.assertEqual(shape_env.guards, [])
+
     def test_strides_from_dim_order(self) -> None:
         sizes = []
         dim_order = []