Support metadata_prop merge and version 25 in version converter (microsoft#2782)

Fix pytorch/pytorch#172784

---

This pull request adds support for ONNX opset version 25 in the version
converter and introduces a new mechanism to copy node metadata during
version conversions. It also includes comprehensive tests to ensure that
metadata is properly transferred to new or replacement nodes created by
adapters during the conversion process.

**Version converter improvements:**

* Increased the maximum supported ONNX opset version from 23 to 25 in
`SUPPORTED_MAX_ONNX_OPSET` within
`onnxscript/version_converter/_version_converter.py`.
* Integrated a new `metadata_merger` utility and implemented a default
metadata merger to ensure node metadata is copied during version
conversion. Metadata is now merged from original nodes to all
replacement nodes in the conversion process.
[[1]](diffhunk://#diff-b6c70f90bafaee79b30e43c90bc0fd5192fb3de7ccc4cf9d48a209798dd775faR239-R244)
[[2]](diffhunk://#diff-b6c70f90bafaee79b30e43c90bc0fd5192fb3de7ccc4cf9d48a209798dd775faR303)

**Testing and validation:**

* Added a new `VersionConverterMetadataMergeTest` class in
`onnxscript/version_converter/_version_converter_test.py` to verify that
metadata is copied correctly to replacement nodes and to all nodes
created by adapters during conversion.
* Updated the test suite to reflect the new maximum supported opset
version and to ensure that conversion beyond version 25 is marked as
expected to fail for future-proofing.

---------

Co-authored-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: titaiwangms

onnxscript/version_converter/_version_converter.py

@@ -12,12 +12,13 @@
 import onnx_ir.convenience as ir_convenience
 
 import onnxscript.ir._tape as _tape
+import onnxscript.utils.metadata_merger as metadata_merger
 from onnxscript import ir
 
 logger = logging.getLogger(__name__)
 
 
-SUPPORTED_MAX_ONNX_OPSET = 23
+SUPPORTED_MAX_ONNX_OPSET = 25
 SUPPORTED_MIN_ONNX_OPSET = 18
 
 
@@ -238,6 +239,12 @@ def groupnormalization_20_21(node: ir.Node, op):
 class _VersionConverter:
     def __init__(self, target_version: int):
         self._target_version = target_version
+        # Default metadata merger: no merging should be needed; keep the first value.
+        self._default_metadata_merger: metadata_merger.MetadataMerger = (
+            metadata_merger.MetadataMerger(
+                dict(),
+            )
+        )
 
     def process_node(
         self, node: ir.Node, from_version: int, up_conversion: bool = True
@@ -293,6 +300,7 @@ def visit_node(
             for new_node in replacement.new_nodes:
                 # TODO: control-flow
                 new_node.version = to_version
+            self._default_metadata_merger.copy_merged_metadata([node], replacement.new_nodes)
             self.replace_node(node, replacement, root)
 
     def visit_graph(self, graph: ir.Graph) -> None:

onnxscript/version_converter/_version_converter_test.py

@@ -296,12 +296,91 @@ def test_version_groupnorm_no_bias(self):
         self.assertEqual(model.graph.node(0).version, 20)
 
 
-class VersionConverter23to24Test(unittest.TestCase):
-    @pytest.mark.xfail(strict=True, reason="Version upgrade beyond 23 not yet supported.")
+class VersionConverterMetadataMergeTest(unittest.TestCase):
+    def test_metadata_is_copied_on_version_conversion(self):
+        """Test that metadata is copied from original node to replacement nodes during version conversion."""
+        model = ir.from_onnx_text(
+            """
+            <ir_version: 7, opset_import: [ "" : 18]>
+            agraph (float[4, 512, 512] input_x) => (float[4, 257, 64, 2] output)
+            {
+                shape_a = Constant<value: tensor = int64[5] {1, 4, 512, 512, 1}>()
+                reshape_x = Reshape (input_x, shape_a)
+                dft = DFT <axis = 2, onesided = 1> (reshape_x)
+                shape_c = Constant<value: tensor = int64[4] {4, 257, 64, 2}>()
+                output = Reshape (dft, shape_c)
+            }
+        """
+        )
+        # Find the DFT node and add metadata to it
+        dft_node = model.graph.node(2)
+        self.assertEqual(dft_node.op_type, "DFT")
+        dft_node.metadata_props["test_key"] = "test_value"
+        dft_node.metadata_props["another_key"] = "another_value"
+
+        target_version = 25
+        version_converter.convert_version(model, target_version=target_version)
+        self.assertEqual(model.opset_imports[""], target_version)
+
+        # After conversion, DFT adapter adds a Constant node for axis and the DFT node is replaced
+        # The replacement DFT node should have the metadata copied
+        new_dft_node = model.graph.node(3)
+        self.assertEqual(new_dft_node.op_type, "DFT")
+        self.assertEqual(new_dft_node.version, 25)
+
+        # Verify metadata was copied to the new DFT node
+        self.assertEqual(new_dft_node.metadata_props.get("test_key"), "test_value")
+        self.assertEqual(new_dft_node.metadata_props.get("another_key"), "another_value")
+
+    def test_metadata_is_copied_to_multiple_replacement_nodes(self):
+        """Test that metadata is copied to all replacement nodes when an adapter creates multiple nodes."""
+        model = ir.from_onnx_text(
+            """
+            <ir_version: 7, opset_import: [ "" : 18]>
+            agraph (float[1, 4, 512, 512] input_x, float[2] scale, float[2] bias) => (float[4, 512, 512] output)
+            {
+                groupnorm = GroupNormalization <num_groups = 2> (input_x, scale, bias)
+                shape_c = Constant<value: tensor = int64[4] {4, 512, 512}>()
+                output = Reshape (groupnorm, shape_c)
+            }
+        """
+        )
+        # Find the GroupNormalization node and add metadata to it
+        groupnorm_node = model.graph.node(0)
+        self.assertEqual(groupnorm_node.op_type, "GroupNormalization")
+        groupnorm_node.metadata_props["source"] = "original_groupnorm"
+
+        target_version = 21
+        version_converter.convert_version(model, target_version=target_version)
+        self.assertEqual(model.opset_imports[""], target_version)
+
+        # GroupNormalization adapter creates multiple nodes (Reshape, Expand, etc.)
+        # Verify that metadata was copied to the new nodes created by the adapter
+        new_groupnorm_node = model.graph.node(9)
+        self.assertEqual(new_groupnorm_node.op_type, "GroupNormalization")
+        self.assertEqual(new_groupnorm_node.version, 21)
+
+        # Verify metadata was copied to the new GroupNormalization node
+        self.assertEqual(new_groupnorm_node.metadata_props.get("source"), "original_groupnorm")
+
+        # Also check that intermediate nodes created by the adapter received the metadata
+        # The adapter creates Reshape, Expand, Reshape nodes for scale and bias
+        for i in range(9):
+            node = model.graph.node(i)
+            if node.version == 21 and node.op_type in ("Reshape", "Expand", "Constant"):
+                self.assertEqual(
+                    node.metadata_props.get("source"),
+                    "original_groupnorm",
+                    f"Node {i} ({node.op_type}) should have metadata copied",
+                )
+
+
+class VersionConverter25to26Test(unittest.TestCase):
+    @pytest.mark.xfail(strict=True, reason="Version upgrade beyond 25 not yet supported.")
     def test_version_convert_compatible(self):
         model = ir.from_onnx_text(
             """
-            <ir_version: 7, opset_import: [ "" : 23]>
+            <ir_version: 7, opset_import: [ "" : 25]>
             agraph (float[1, 4, 512, 512] input_x, float[1, 4, 512, 64] input_y) => (float[1, 4, 512, 64] output)
             {
                 shape_a = Constant<value: tensor = int64[3] {4, 512, 512}>()
@@ -314,7 +393,7 @@ def test_version_convert_compatible(self):
             }
         """
         )
-        target_version = 24
+        target_version = 26
         version_converter.convert_version(model, target_version=target_version)