Channels last: preserve metadata across op replacement

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