Channels last: preserve metadata across op replacement (pytorch#21961)

Preserve complete operator metadata when replacing ATen operators with
channels-last dialect counterparts. ExportPass retracing recomputes
layout-dependent shape fields while source attribution, quantization
annotations, and custom backend metadata survive.

Document that callers must reject or remap semantic metadata tied to
dimensions changed by the replacement specification, such as per-channel
quantization axes.

pytest -q
backends/transforms/test/test_replace_ops_with_channels_last_variants.py
(23 passed).

Authored with Codex.

Author: rascani

backends/transforms/replace_ops_with_channels_last_variants.py

@@ -133,6 +133,11 @@ class ReplaceOpsWithChannelsLastVariants(ExportPass):
 
     By default, all currently implemented channels_last dialect ops are replaced.
     Pass a custom op_map to restrict or extend the set of replacements.
+
+    Metadata from each replaced operator is preserved so provenance and backend
+    annotations survive the rewrite. ExportPass recomputes shape metadata after
+    retracing. Callers must reject or remap semantic metadata tied to dimensions
+    changed by ``input_indices`` or ``output_indices``, such as per-channel axes.
     """
 
     def __init__(
@@ -229,7 +234,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                     args=tuple(args),
                     kwargs=node.kwargs,
                 )
-                nhwc_node.meta = {}
+                nhwc_node.meta = dict(node.meta)
 
                 users = list(node.users)
                 if all(

backends/transforms/test/test_replace_ops_with_channels_last_variants.py

@@ -132,6 +132,26 @@ def forward(self, x):
 
 
 class TestReplaceOpsWithChannelsLastVariants:
+    def test_preserves_metadata_and_recomputes_shape(self):
+        ep = _export_to_edge(Conv2dModule(), (torch.randn(1, 4, 8, 8),))
+        conv = _find_nodes(ep.graph_module, exir_ops.edge.aten.convolution.default)[0]
+        metadata = {
+            "debug_handle": 1234,
+            "from_node": [("source", "convolution")],
+            "input_qparams": {0: "input"},
+            "output_qparams": {0: "output"},
+        }
+        conv.meta.update(metadata)
+
+        result = ReplaceOpsWithChannelsLastVariants(ep)(ep.graph_module)
+        replaced = _find_nodes(
+            result.graph_module, exir_ops.edge.channels_last.convolution.default
+        )[0]
+
+        for key, value in metadata.items():
+            assert replaced.meta[key] == value
+        assert tuple(replaced.meta["val"].shape) == (1, 8, 8, 4)
+
     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