diff --git a/backends/transforms/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py index 5540fe55f22..5e8a4612163 100644 --- a/backends/transforms/replace_ops_with_channels_last_variants.py +++ b/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( 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 index e131463320e..3c8fdd8309f 100644 --- a/backends/transforms/test/test_replace_ops_with_channels_last_variants.py +++ b/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