diff --git a/backends/transforms/channels_last_ops.py b/backends/transforms/channels_last_ops.py
index 54a897c7b80..8a2d7401bb3 100644
--- a/backends/transforms/channels_last_ops.py
+++ b/backends/transforms/channels_last_ops.py
@@ -40,8 +40,17 @@ def _conv(
 
 
 def _avg_pool2d(
-    input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override
+    input,
+    kernel_size,
+    stride=None,
+    padding=0,
+    ceil_mode=False,
+    count_include_pad=True,
+    divisor_override=None,
 ):
+    if stride is None:
+        stride = []  # Default value of avg_pool2d.
+
     nchw = input.permute(0, 3, 1, 2)
     out = torch.ops.aten.avg_pool2d(
         nchw,
@@ -75,7 +84,12 @@ def _upsample_nearest2d(input, output_size, scale_factors):
     return out.permute(0, 2, 3, 1).contiguous()
 
 
-def _max_pool2d_with_indices(input, kernel_size, stride, padding, dilation, ceil_mode):
+def _max_pool2d_with_indices(
+    input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False
+):
+    if stride is None:
+        stride = []  # Default value of max_pool2d_with_indices.
+
     nchw = input.permute(0, 3, 1, 2)
     values, indices = torch.ops.aten.max_pool2d_with_indices(
         nchw, kernel_size, stride, padding, dilation, ceil_mode
@@ -106,8 +120,8 @@ def _permute_copy(input, dims):
 register_fake("channels_last::convolution", _conv, lib=lib)
 
 lib.define(
-    "avg_pool2d(Tensor input, int[2] kernel_size, int[2] stride, int[2] padding, "
-    "bool ceil_mode, bool count_include_pad, int? divisor_override) -> Tensor"
+    "avg_pool2d(Tensor input, int[2] kernel_size, int[2] stride=[], int[2] padding=0, "
+    "bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> Tensor"
 )
 lib.impl("avg_pool2d", _avg_pool2d, "CompositeExplicitAutograd")
 register_fake("channels_last::avg_pool2d", _avg_pool2d, lib=lib)
@@ -131,8 +145,8 @@ def _permute_copy(input, dims):
 register_fake("channels_last::upsample_nearest2d", _upsample_nearest2d, lib=lib)
 
 lib.define(
-    "max_pool2d_with_indices(Tensor input, int[2] kernel_size, int[2] stride, "
-    "int[2] padding, int[2] dilation, bool ceil_mode) -> (Tensor, Tensor)"
+    "max_pool2d_with_indices(Tensor input, int[2] kernel_size, int[2] stride=[], "
+    "int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor)"
 )
 lib.impl(
     "max_pool2d_with_indices", _max_pool2d_with_indices, "CompositeExplicitAutograd"
diff --git a/backends/transforms/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py
new file mode 100644
index 00000000000..63a1ca5a014
--- /dev/null
+++ b/backends/transforms/replace_ops_with_channels_last_variants.py
@@ -0,0 +1,269 @@
+# Copyright 2026 NXP
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import operator
+from dataclasses import dataclass
+from typing import Callable
+
+import executorch.backends.transforms.channels_last_ops  # noqa: F401
+
+import torch
+
+from executorch.exir import ExportedProgram
+from executorch.exir.dialects._ops import ops as exir_ops
+
+from executorch.exir.pass_base import ExportPass
+from torch.fx.node import Target
+from torch.fx.passes.infra.pass_manager import PassResult
+
+_NCHW_TO_NHWC_PERM: list[int] = [0, 2, 3, 1]
+_NHWC_TO_NCHW_PERM: list[int] = [0, 3, 1, 2]
+
+
+FilterFnType = Callable[[torch.fx.Node], bool]
+
+
+def _requires_rank(allowed_ranks: list[int]) -> FilterFnType:
+    def _inner(node) -> bool:
+        val = node.meta["val"]
+        val = val[0] if isinstance(val, (list, tuple)) else val
+        return val.dim() in allowed_ranks
+
+    return _inner
+
+
+def _has_2d_convolution_kernel() -> FilterFnType:
+    def _inner(node) -> bool:
+        # Extract kernel size from the weights
+        try:
+            w = node.args[1]
+            return w.meta["val"].dim() == 4
+        except AttributeError:
+            return False
+
+    return _inner
+
+
+def _and(*filter_fns: FilterFnType) -> FilterFnType:
+    def _inner(*args) -> bool:
+        return all(fn(*args) for fn in filter_fns)
+
+    return _inner
+
+
+@dataclass
+class ChannelsLastOpSpec:
+    """Specification for replacing a contiguous op with its channels-last counterpart."""
+
+    # The channels_last dialect target to replace the contiguous op with.
+    target: Target
+
+    # Positional arg indices of tensor inputs that should be permuted NCHW→NHWC.
+    input_indices: list[int]
+
+    # Indices of the outputs that should be permuted NCHW→NHWC.
+    output_indices: list[int]
+
+    # If provided, this function must return True for a node to be replaced.
+    filter_fn: FilterFnType | None = None
+
+
+_DEFAULT_OP_MAP: dict[Target, ChannelsLastOpSpec] = {
+    exir_ops.edge.aten._adaptive_avg_pool2d.default: ChannelsLastOpSpec(
+        target=exir_ops.edge.channels_last.adaptive_avg_pool2d.default,
+        input_indices=[0],
+        output_indices=[0],
+        filter_fn=_requires_rank([3, 4]),
+    ),
+    exir_ops.edge.aten.avg_pool2d.default: ChannelsLastOpSpec(
+        target=exir_ops.edge.channels_last.avg_pool2d.default,
+        input_indices=[0],
+        output_indices=[0],
+        filter_fn=_requires_rank([3, 4]),
+    ),
+    exir_ops.edge.aten.convolution.default: ChannelsLastOpSpec(
+        target=exir_ops.edge.channels_last.convolution.default,
+        input_indices=[0],
+        output_indices=[0],
+        filter_fn=_and(_requires_rank([3, 4]), _has_2d_convolution_kernel()),
+    ),
+    exir_ops.edge.aten.grid_sampler_2d.default: ChannelsLastOpSpec(
+        target=exir_ops.edge.channels_last.grid_sampler_2d.default,
+        input_indices=[0],
+        output_indices=[0],
+        filter_fn=_requires_rank([4]),
+    ),
+    exir_ops.edge.aten.max_pool2d_with_indices.default: ChannelsLastOpSpec(
+        target=exir_ops.edge.channels_last.max_pool2d_with_indices.default,
+        input_indices=[0],
+        output_indices=[0, 1],
+        filter_fn=_requires_rank([3, 4]),
+    ),
+    exir_ops.edge.aten.upsample_bilinear2d.vec: ChannelsLastOpSpec(
+        target=exir_ops.edge.channels_last.upsample_bilinear2d.default,
+        input_indices=[0],
+        output_indices=[0],
+        filter_fn=_requires_rank([4]),
+    ),
+    exir_ops.edge.aten.upsample_nearest2d.vec: ChannelsLastOpSpec(
+        target=exir_ops.edge.channels_last.upsample_nearest2d.default,
+        input_indices=[0],
+        output_indices=[0],
+        filter_fn=_requires_rank([4]),
+    ),
+}
+
+
+class ReplaceOpsWithChannelsLastVariants(ExportPass):
+    """
+    Replaces contiguous (NCHW) edge-dialect ops with their channels-last (NHWC)
+    equivalents from the channels_last dialect, inserting permute_copy ops on
+    the specified inputs and on the output so that the rest of the graph remains
+    in contiguous layout.
+
+    By default, all currently implemented channels_last dialect ops are replaced.
+    Pass a custom op_map to restrict or extend the set of replacements.
+    """
+
+    def __init__(
+        self,
+        exported_program: ExportedProgram,
+        op_map: dict[Target, ChannelsLastOpSpec] | None = None,
+    ) -> None:
+        super().__init__()
+        self.exported_program = exported_program
+        self.op_map: dict[Target, ChannelsLastOpSpec] = (
+            op_map if op_map is not None else dict(_DEFAULT_OP_MAP)
+        )
+
+    @staticmethod
+    def _permute_node_input(
+        graph: torch.fx.Graph, node_input: torch.fx.Node, implicit_batch: bool
+    ) -> torch.fx.Node:
+        if implicit_batch:
+            # Explicitly add batch size of `1`.
+            node_input = graph.create_node(
+                "call_function",
+                target=exir_ops.edge.aten.unsqueeze_copy.default,
+                args=(node_input, 0),
+            )
+            node_input.meta = {}
+
+        res = graph.create_node(
+            "call_function",
+            target=exir_ops.edge.channels_last.permute_copy.default,
+            args=(node_input, _NCHW_TO_NHWC_PERM),
+        )
+        res.meta = {}
+
+        return res
+
+    @staticmethod
+    def _permute_node_output(
+        graph: torch.fx.Graph,
+        node_output: torch.fx.Node,
+        original_node_output: torch.fx.Node,
+        implicit_batch: bool,
+    ):
+        output = graph.create_node(
+            "call_function",
+            target=exir_ops.edge.channels_last.permute_copy.default,
+            args=(node_output, _NHWC_TO_NCHW_PERM),
+        )
+        output.meta = {}
+
+        if implicit_batch:
+            # Remove the explicitly added batch size of `1`.
+            output = graph.create_node(
+                "call_function",
+                target=exir_ops.edge.aten.squeeze_copy.dims,
+                args=(output, [0]),
+            )
+            output.meta = {}
+
+        original_node_output.replace_all_uses_with(output)
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        modified = False
+        graph = graph_module.graph
+
+        for node in list(graph.nodes):
+            if node.op != "call_function":
+                continue
+            if (spec := self.op_map.get(node.target)) is None:
+                continue
+            val = node.meta["val"]
+            val = val[0] if isinstance(val, (list, tuple)) else val
+            contiguous_dim_order = tuple(range(val.dim()))
+            if val.dim_order() != contiguous_dim_order:
+                continue
+            if spec.filter_fn is not None and not spec.filter_fn(node):
+                continue
+
+            # In case of implicit batch size, insert also `unsqueeze_copy.default` and `squeeze_copy.dims` operators.
+            #  With `convolution`, this already happens during lowering to edge. But it doesn't happen for example with
+            #  `avg_pool2d`. Therefore, we do it manually here to achieve consistency between the operators.
+            implicit_batch = val.dim() == 3
+
+            args = list(node.args)
+
+            with graph.inserting_before(node):
+                for idx in spec.input_indices:
+                    args[idx] = self._permute_node_input(
+                        graph, args[idx], implicit_batch
+                    )
+
+                nhwc_node = graph.create_node(
+                    "call_function",
+                    target=spec.target,
+                    args=tuple(args),
+                    kwargs=node.kwargs,
+                )
+                nhwc_node.meta = {}
+
+                users = list(node.users)
+                if all(
+                    u.op == "call_function" and u.target == operator.getitem
+                    for u in users
+                ):
+                    # `node` produces multiple outputs which are extracted by following `getitem` nodes.
+                    for idx in spec.output_indices:
+                        old_getitem_nodes = [u for u in users if u.args[1] == idx]
+                        if not old_getitem_nodes:
+                            continue  # The output is not used.
+
+                        # Add a new `getitem` node.
+                        new_getitem_node = graph.create_node(
+                            "call_function",
+                            target=operator.getitem,
+                            args=(nhwc_node, idx),
+                        )
+                        new_getitem_node.meta = {}
+
+                        self._permute_node_output(
+                            graph,
+                            new_getitem_node,
+                            old_getitem_nodes[0],
+                            implicit_batch,
+                        )
+
+                        # Remove the old `getitem` node.
+                        graph.erase_node(old_getitem_nodes[0])
+                else:
+                    # Regular node with a single output.
+                    assert spec.output_indices == [
+                        0
+                    ], f"Incorrect use of `output_indices` for op `{spec.target}`."
+                    self._permute_node_output(graph, nhwc_node, node, implicit_batch)
+
+            graph.erase_node(node)
+            modified = True
+
+        if modified:
+            graph.eliminate_dead_code()
+            graph_module.recompile()
+            graph_module = super().call(graph_module).graph_module
+
+        return PassResult(graph_module, modified)
diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl
index 253a7b57bc9..61297e11cb4 100644
--- a/backends/transforms/targets.bzl
+++ b/backends/transforms/targets.bzl
@@ -463,3 +463,33 @@ def define_common_targets():
             ":replace_nop_transpose_or_permute_with_view",
         ],
     )
+
+    runtime.python_library(
+        name = "replace_ops_with_channels_last_variants",
+        srcs = [
+            "replace_ops_with_channels_last_variants.py",
+        ],
+        visibility = [
+            "//executorch/backends/...",
+        ],
+        deps = [
+            "//caffe2:torch",
+            ":channels_last_ops",
+            "//executorch/exir:pass_base",
+            "//executorch/exir:lib",
+        ],
+    )
+
+    runtime.python_test(
+        name = "test_replace_ops_with_channels_last_variants",
+        srcs = [
+            "test/test_replace_ops_with_channels_last_variants.py",
+        ],
+        deps = [
+            "//caffe2:torch",
+            ":channels_last_ops",
+            ":replace_ops_with_channels_last_variants",
+            "//executorch/exir:lib",
+            "fbsource//third-party/pypi/pytest:pytest",
+        ],
+    )
diff --git a/backends/transforms/test/test_replace_ops_with_channels_last_variants.py b/backends/transforms/test/test_replace_ops_with_channels_last_variants.py
new file mode 100644
index 00000000000..6f1807205a4
--- /dev/null
+++ b/backends/transforms/test/test_replace_ops_with_channels_last_variants.py
@@ -0,0 +1,505 @@
+# Copyright 2026 NXP
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import operator
+
+import executorch.backends.transforms.channels_last_ops  # noqa: F401
+import pytest
+import torch
+from executorch.backends.transforms.replace_ops_with_channels_last_variants import (
+    _NCHW_TO_NHWC_PERM,
+    _NHWC_TO_NCHW_PERM,
+    ChannelsLastOpSpec,
+    ReplaceOpsWithChannelsLastVariants,
+)
+from executorch.exir import to_edge
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export import ExportedProgram
+from torch.fx import GraphModule
+from torch.fx.node import Target
+
+
+# ── helpers ────────────────────────────────────────────────────────────────────
+
+
+def _export_to_edge(module: torch.nn.Module, inputs: tuple) -> ExportedProgram:
+    ep = torch.export.export(module.eval(), inputs)
+    return to_edge(ep).exported_program()
+
+
+def _run_pass(ep: ExportedProgram) -> tuple[GraphModule, bool]:
+    result = ReplaceOpsWithChannelsLastVariants(ep)(ep.graph_module)
+    return result.graph_module, result.modified
+
+
+def _find_nodes(gm: GraphModule, target: Target) -> list[torch.fx.Node]:
+    return [n for n in gm.graph.nodes if n.op == "call_function" and n.target == target]
+
+
+def _count(gm: GraphModule, target: Target) -> int:
+    return len(_find_nodes(gm, target))
+
+
+# ── modules ────────────────────────────────────────────────────────────────────
+
+
+class Conv2dModule(torch.nn.Module):
+    def __init__(self, groups: int = 1, bias: bool = True):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(
+            4, 4, kernel_size=3, padding=1, groups=groups, bias=bias
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+class AvgPool2dModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+
+
+class AdaptiveAvgPool2dModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.adaptive_avg_pool2d(x, (2, 2))
+
+
+class Conv2dAvgPool2dModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(4, 4, kernel_size=3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.avg_pool2d(self.conv(x), kernel_size=2, stride=2)
+
+
+class GridSampler2DModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.grid_sample(x, torch.ones(1, 2, 2, 2))
+
+
+class Conv1dModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv1d(4, 4, kernel_size=3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+class LinearModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc = torch.nn.Linear(8, 4)
+
+    def forward(self, x):
+        return self.fc(x)
+
+
+class MaxPool2DModule(torch.nn.Module):
+    def __init__(self, kernel_size: int | tuple[int, ...] = 3, **kwargs):
+        super().__init__()
+        self.max_pool2d = torch.nn.MaxPool2d(kernel_size, **kwargs)
+
+    def forward(self, x):
+        return self.max_pool2d(x)
+
+
+class UpsampleBilinearModule(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.upsample = torch.nn.Upsample(scale_factor=2, mode="bilinear")
+
+    def forward(self, x):
+        return self.upsample(x)
+
+
+class UpsampleNearestModule(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.upsample = torch.nn.Upsample(scale_factor=2, mode="nearest")
+
+    def forward(self, x):
+        return self.upsample(x)
+
+
+# ── tests ──────────────────────────────────────────────────────────────────────
+
+
+class TestReplaceOpsWithChannelsLastVariants:
+
+    def test_conv2d(self):
+        ep = _export_to_edge(Conv2dModule(bias=True), (torch.randn(1, 4, 8, 8),))
+        assert _count(ep.graph_module, exir_ops.edge.aten.convolution.default) == 1
+
+        gm, modified = _run_pass(ep)
+
+        assert modified
+        assert _count(gm, exir_ops.edge.aten.convolution.default) == 0
+        assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 1
+        assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2
+
+        dialect_node = _find_nodes(gm, exir_ops.edge.channels_last.convolution.default)[
+            0
+        ]
+
+        # Input (arg 0) must be wrapped in an NCHW→NHWC permute.
+        input_permute = dialect_node.args[0]
+        assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM
+
+        # Weight (arg 1) and bias (arg 2) must not be permuted.
+        for arg in (dialect_node.args[1], dialect_node.args[2]):
+            if isinstance(arg, torch.fx.Node):
+                assert arg.target != exir_ops.edge.channels_last.permute_copy.default
+
+        # Output must be wrapped in an NHWC→NCHW permute.
+        users = list(dialect_node.users)
+        assert len(users) == 1
+        output_permute = users[0]
+        assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM
+
+    def test_conv2d__implicit_batch(self):
+        input_shape = (4, 8, 8)  # Use implicit batch size of `1`.
+        ep = _export_to_edge(Conv2dModule(bias=True), (torch.randn(*input_shape),))
+        assert _count(ep.graph_module, exir_ops.edge.aten.convolution.default) == 1
+
+        gm, modified = _run_pass(ep)
+
+        assert modified
+        assert _count(gm, exir_ops.edge.aten.convolution.default) == 0
+        assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 1
+        assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2
+
+        # These ops are added automatically during lowering to edge.
+        assert _count(gm, exir_ops.edge.aten.unsqueeze_copy.default) == 1
+        assert _count(gm, exir_ops.edge.aten.squeeze_copy.dims) == 1
+
+        dialect_node = _find_nodes(gm, exir_ops.edge.channels_last.convolution.default)[
+            0
+        ]
+
+        # Input (arg 0) must be wrapped in an NCHW→NHWC permute.
+        input_permute = dialect_node.args[0]
+        assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM
+        assert input_permute.args[0].target == exir_ops.edge.aten.unsqueeze_copy.default
+
+        # Weight (arg 1) and bias (arg 2) must not be permuted.
+        for arg in (dialect_node.args[1], dialect_node.args[2]):
+            if isinstance(arg, torch.fx.Node):
+                assert arg.target != exir_ops.edge.channels_last.permute_copy.default
+
+        # Output must be wrapped in an NHWC→NCHW permute.
+        users = list(dialect_node.users)
+        assert len(users) == 1
+        output_permute = users[0]
+        assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM
+        assert (
+            list(output_permute.users)[0].target == exir_ops.edge.aten.squeeze_copy.dims
+        )
+
+    def test_depthwise_conv2d(self):
+        ep = _export_to_edge(Conv2dModule(groups=4), (torch.randn(1, 4, 8, 8),))
+        gm, modified = _run_pass(ep)
+
+        assert modified
+        assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 1
+        assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2
+
+    def test_max_pool2d(self):
+        ep = _export_to_edge(MaxPool2DModule(), (torch.randn(1, 4, 8, 8),))
+        assert (
+            _count(ep.graph_module, exir_ops.edge.aten.max_pool2d_with_indices.default)
+            == 1
+        )
+
+        gm, modified = _run_pass(ep)
+
+        assert modified
+        assert _count(gm, exir_ops.edge.aten.max_pool2d_with_indices.default) == 0
+        assert (
+            _count(gm, exir_ops.edge.channels_last.max_pool2d_with_indices.default) == 1
+        )
+        assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2
+        assert _count(gm, operator.getitem) == 1
+
+        dialect_node = _find_nodes(
+            gm, exir_ops.edge.channels_last.max_pool2d_with_indices.default
+        )[0]
+
+        # Input (arg 0) must be wrapped in an NCHW→NHWC permute.
+        input_permute = dialect_node.args[0]
+        assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM
+
+        users = list(dialect_node.users)
+        assert len(users) == 1
+        output_getitem = users[0]
+        assert output_getitem.target == operator.getitem
+
+        # Output must be wrapped in an NHWC→NCHW permute.
+        getitem_users = list(output_getitem.users)
+        assert len(getitem_users) == 1
+        output_permute = getitem_users[0]
+        assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM
+
+    def test_max_pool2d__implicit_batch(self):
+        input_shape = (4, 8, 8)  # Use implicit batch size of `1`.
+        ep = _export_to_edge(MaxPool2DModule(), (torch.randn(*input_shape),))
+        assert (
+            _count(ep.graph_module, exir_ops.edge.aten.max_pool2d_with_indices.default)
+            == 1
+        )
+
+        gm, modified = _run_pass(ep)
+
+        assert modified
+        assert _count(gm, exir_ops.edge.aten.max_pool2d_with_indices.default) == 0
+        assert (
+            _count(gm, exir_ops.edge.channels_last.max_pool2d_with_indices.default) == 1
+        )
+        assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2
+        assert _count(gm, exir_ops.edge.aten.unsqueeze_copy.default) == 1
+        assert _count(gm, exir_ops.edge.aten.squeeze_copy.dims) == 1
+
+        dialect_node = _find_nodes(
+            gm, exir_ops.edge.channels_last.max_pool2d_with_indices.default
+        )[0]
+
+        # Input (arg 0) must be wrapped in an NCHW→NHWC permute + an unsqueeze.
+        input_permute = dialect_node.args[0]
+        assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM
+        assert input_permute.args[0].target == exir_ops.edge.aten.unsqueeze_copy.default
+
+        users = list(dialect_node.users)
+        assert len(users) == 1
+        output_getitem = users[0]
+        assert output_getitem.target == operator.getitem
+
+        # Output must be wrapped in an NHWC→NCHW permute + a squeeze.
+        getitem_users = list(output_getitem.users)
+        assert len(getitem_users) == 1
+        output_permute = getitem_users[0]
+        assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM
+        assert (
+            list(output_permute.users)[0].target == exir_ops.edge.aten.squeeze_copy.dims
+        )
+
+    @pytest.mark.parametrize(
+        "module_cls, aten_op, channels_last_op",
+        [
+            pytest.param(
+                AvgPool2dModule,
+                exir_ops.edge.aten.avg_pool2d.default,
+                exir_ops.edge.channels_last.avg_pool2d.default,
+                id="AvgPool2D",
+            ),
+            pytest.param(
+                AdaptiveAvgPool2dModule,
+                exir_ops.edge.aten._adaptive_avg_pool2d.default,
+                exir_ops.edge.channels_last.adaptive_avg_pool2d.default,
+                id="AdaptiveAvgPool2D",
+            ),
+            pytest.param(
+                GridSampler2DModule,
+                exir_ops.edge.aten.grid_sampler_2d.default,
+                exir_ops.edge.channels_last.grid_sampler_2d.default,
+                id="GridSampler2D",
+            ),
+            pytest.param(
+                UpsampleBilinearModule,
+                exir_ops.edge.aten.upsample_bilinear2d.vec,
+                exir_ops.edge.channels_last.upsample_bilinear2d.default,
+                id="UpsampleBilinear2D",
+            ),
+            pytest.param(
+                UpsampleNearestModule,
+                exir_ops.edge.aten.upsample_nearest2d.vec,
+                exir_ops.edge.channels_last.upsample_nearest2d.default,
+                id="UpsampleNearest2D",
+            ),
+        ],
+    )
+    def test_ops_without_weights(self, module_cls, aten_op, channels_last_op):
+        ep = _export_to_edge(module_cls(), (torch.randn(1, 4, 8, 8),))
+        assert _count(ep.graph_module, aten_op) == 1
+
+        gm, modified = _run_pass(ep)
+
+        assert modified
+        assert _count(gm, aten_op) == 0
+        assert _count(gm, channels_last_op) == 1
+        assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2
+
+        dialect_node = _find_nodes(gm, channels_last_op)[0]
+
+        input_permute = dialect_node.args[0]
+        assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM
+
+        users = list(dialect_node.users)
+        assert len(users) == 1
+        output_permute = users[0]
+        assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM
+
+    @pytest.mark.parametrize(
+        "module_cls, aten_op, channels_last_op",
+        [
+            # GridSampler2D, UpsampleBilinear2D and UpsampleNearest2D don't support implicit batch.
+            pytest.param(
+                AvgPool2dModule,
+                exir_ops.edge.aten.avg_pool2d.default,
+                exir_ops.edge.channels_last.avg_pool2d.default,
+                id="AvgPool2D",
+            ),
+            pytest.param(
+                AdaptiveAvgPool2dModule,
+                exir_ops.edge.aten._adaptive_avg_pool2d.default,
+                exir_ops.edge.channels_last.adaptive_avg_pool2d.default,
+                id="AdaptiveAvgPool2D",
+            ),
+        ],
+    )
+    def test_ops_without_weights__implicit_batch(
+        self, module_cls, aten_op, channels_last_op
+    ):
+        input_shape = (4, 8, 8)  # Use implicit batch size of `1`.
+        ep = _export_to_edge(module_cls(), (torch.randn(*input_shape),))
+        assert _count(ep.graph_module, aten_op) == 1
+
+        gm, modified = _run_pass(ep)
+
+        assert modified
+        assert _count(gm, aten_op) == 0
+        assert _count(gm, channels_last_op) == 1
+        assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2
+        assert _count(gm, exir_ops.edge.aten.unsqueeze_copy.default) == 1
+        assert _count(gm, exir_ops.edge.aten.squeeze_copy.dims) == 1
+
+        dialect_node = _find_nodes(gm, channels_last_op)[0]
+
+        input_permute = dialect_node.args[0]
+        assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM
+        assert input_permute.args[0].target == exir_ops.edge.aten.unsqueeze_copy.default
+
+        users = list(dialect_node.users)
+        assert len(users) == 1
+        output_permute = users[0]
+        assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default
+        assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM
+        assert (
+            list(output_permute.users)[0].target == exir_ops.edge.aten.squeeze_copy.dims
+        )
+
+    def test_conv2d_and_avg_pool2d(self):
+        # 2 permutes per op. The redundant adjacent pair (NHWC→NCHW then NCHW→NHWC) between the two ops will be
+        #  optimized out by a different pass.
+        ep = _export_to_edge(Conv2dAvgPool2dModule(), (torch.randn(1, 4, 8, 8),))
+        gm, modified = _run_pass(ep)
+
+        assert modified
+        assert _count(gm, exir_ops.edge.aten.convolution.default) == 0
+        assert _count(gm, exir_ops.edge.aten.avg_pool2d.default) == 0
+        assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 1
+        assert _count(gm, exir_ops.edge.channels_last.avg_pool2d.default) == 1
+        assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 4
+
+    def test_modified_false_when_no_matching_ops(self):
+        ep = _export_to_edge(LinearModule(), (torch.randn(2, 8),))
+        _, modified = _run_pass(ep)
+        assert not modified
+
+    def test_empty_op_map_leaves_graph_unchanged(self):
+        ep = _export_to_edge(Conv2dModule(), (torch.randn(1, 4, 8, 8),))
+
+        result = ReplaceOpsWithChannelsLastVariants(ep, op_map={})(ep.graph_module)
+
+        assert not result.modified
+        assert _count(result.graph_module, exir_ops.edge.aten.convolution.default) == 1
+        assert (
+            _count(result.graph_module, exir_ops.edge.channels_last.convolution.default)
+            == 0
+        )
+
+    def test_custom_op_map_only_replaces_specified_op(self):
+        ep = _export_to_edge(Conv2dAvgPool2dModule(), (torch.randn(1, 4, 8, 8),))
+
+        custom_map = {
+            exir_ops.edge.aten.avg_pool2d.default: ChannelsLastOpSpec(
+                target=exir_ops.edge.channels_last.avg_pool2d.default,
+                input_indices=[0],
+                output_indices=[0],
+            )
+        }
+        result = ReplaceOpsWithChannelsLastVariants(ep, op_map=custom_map)(
+            ep.graph_module
+        )
+        gm = result.graph_module
+
+        assert result.modified
+        assert _count(gm, exir_ops.edge.aten.convolution.default) == 1
+        assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 0
+        assert _count(gm, exir_ops.edge.aten.avg_pool2d.default) == 0
+        assert _count(gm, exir_ops.edge.channels_last.avg_pool2d.default) == 1
+
+    def test_conv1d__not_replaced(self):
+        ep = _export_to_edge(Conv1dModule(), (torch.randn(1, 4, 16),))
+        gm, modified = _run_pass(ep)
+
+        assert not modified
+        assert _count(gm, exir_ops.edge.aten.convolution.default) == 1
+        assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 0
+
+    def test_convolution__numerical_correctness(self):
+        torch.manual_seed(2)
+        x = torch.randn(1, 4, 8, 8)
+        model = Conv2dModule(bias=False)
+        reference_output = model(x)
+
+        ep = _export_to_edge(model, (x,))
+        gm, modified = _run_pass(ep)
+        assert modified
+        assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 1
+
+        channels_last_output = gm(ep.state_dict["conv.weight"], x)[0]
+        assert torch.allclose(channels_last_output, reference_output)
+
+    def test_avg_pool2d__numerical_correctness(self):
+        torch.manual_seed(2)
+        x = torch.randn(1, 4, 8, 8)
+        model = AvgPool2dModule()
+        reference_output = model(x)
+
+        ep = _export_to_edge(model, (x,))
+        gm, modified = _run_pass(ep)
+        assert modified
+        assert _count(gm, exir_ops.edge.channels_last.avg_pool2d.default) == 1
+
+        channels_last_output = gm(x)[0]
+        assert torch.allclose(channels_last_output, reference_output)
+
+    def test_avg_pool2d__implicit_batch__numerical_correctness(self):
+        torch.manual_seed(2)
+        x = torch.randn(4, 8, 8)
+        model = AvgPool2dModule()
+        reference_output = model(x)
+
+        ep = _export_to_edge(model, (x,))
+        gm, modified = _run_pass(ep)
+        assert modified
+        assert _count(gm, exir_ops.edge.channels_last.avg_pool2d.default) == 1
+
+        channels_last_output = gm(x)[0]
+        assert torch.allclose(channels_last_output, reference_output)