Fix version converter regression (#2799)

This pull request improves the ONNX version converter by enabling direct
opset version conversion for models containing functions, and ensures
that version conversion now correctly applies to function nodes and
their subgraphs. The changes also update the internal APIs to handle
both models and functions, and expand test coverage for these scenarios.

Enhancements to version conversion logic:

* Removed the requirement to inline functions before version conversion,
allowing `_ConvertVersionPass` to process models with functions directly
(`onnxscript/version_converter/__init__.py`).
[[1]](diffhunk://#diff-e2693e79e7e645175e987e44ba3432063a21003bf4aa50119c1a786925f6e042L41-R41)
[[2]](diffhunk://#diff-e2693e79e7e645175e987e44ba3432063a21003bf4aa50119c1a786925f6e042L55-R54)
[[3]](diffhunk://#diff-e2693e79e7e645175e987e44ba3432063a21003bf4aa50119c1a786925f6e042L76-L81)
* Updated the internal API (`_set_onnx_opset_version`) to support both
models and functions, so opset imports are set correctly for functions
as well as models
(`onnxscript/version_converter/_version_converter.py`).

Support for function and subgraph conversion:

* Modified traversal logic to process nodes inside functions and
subgraphs within control flow nodes, ensuring that all relevant nodes
are version-converted
(`onnxscript/version_converter/_version_converter.py`).
[[1]](diffhunk://#diff-b6c70f90bafaee79b30e43c90bc0fd5192fb3de7ccc4cf9d48a209798dd775faL277-R280)
[[2]](diffhunk://#diff-b6c70f90bafaee79b30e43c90bc0fd5192fb3de7ccc4cf9d48a209798dd775faL306-R307)
[[3]](diffhunk://#diff-b6c70f90bafaee79b30e43c90bc0fd5192fb3de7ccc4cf9d48a209798dd775faL324-R324)
[[4]](diffhunk://#diff-b6c70f90bafaee79b30e43c90bc0fd5192fb3de7ccc4cf9d48a209798dd775faL334-R337)

Expanded test coverage:

* Added tests to verify that nodes inside functions and subgraphs within
control flow nodes (e.g., `If` branches) are correctly
version-converted, and that opset imports are updated for functions
(`onnxscript/version_converter/_version_converter_test.py`).

Author: titaiwangms

onnxscript/version_converter/_version_converter.py

@@ -34,11 +34,11 @@ def _get_onnx_opset_version(model: ir.Model) -> int | None:
     return model_version1 or model_version2
 
 
-def _set_onnx_opset_version(model: ir.Model, version: int) -> None:
-    """Set the ONNX opset version imported by the model."""
-    if "ai.onnx" in model.opset_imports:
-        del model.opset_imports["ai.onnx"]
-    model.opset_imports[""] = version
+def _set_onnx_opset_version(model_or_function: ir.Model | ir.Function, version: int) -> None:
+    """Set the ONNX opset version imported by the model or function."""
+    if "ai.onnx" in model_or_function.opset_imports:
+        del model_or_function.opset_imports["ai.onnx"]
+    model_or_function.opset_imports[""] = version
 
 
 class VersionConverterError(RuntimeError):
@@ -334,6 +334,7 @@ def visit_model(self, model: ir.Model) -> None:
         self.visit_graph_or_function(model.graph)
         for function in model.functions.values():
             self.visit_graph_or_function(function)
+            _set_onnx_opset_version(function, self._target_version)
         _set_onnx_opset_version(model, self._target_version)
 
 

onnxscript/version_converter/_version_converter_test.py

@@ -237,8 +237,11 @@ def test_version_convert_function_nodes(self):
         version_converter.convert_version(model, target_version=target_version)
         self.assertEqual(model.opset_imports[""], target_version)
 
-        # Verify that nodes inside the function were version-converted
+        # Verify that the function's opset_imports are updated
         func = model.functions[("pkg.custom", "dft_func", "")]
+        self.assertEqual(func.opset_imports[""], target_version)
+
+        # Verify that nodes inside the function were version-converted
         self.assertEqual(func[0].op_type, "Constant")
         self.assertEqual(func[0].version, 20)
         self.assertEqual(func[1].op_type, "Reshape")
@@ -293,8 +296,12 @@ def test_version_convert_function_with_control_flow_subgraph(self):
         version_converter.convert_version(model, target_version=target_version)
         self.assertEqual(model.opset_imports[""], target_version)
 
-        # Verify nodes inside the function's If node subgraphs were version-converted
+        # Verify that the function's opset_imports are updated
         func = model.functions[("pkg.custom", "conditional_dft", "")]
+        self.assertEqual(func.opset_imports[""], target_version)
+
+        # Verify nodes inside the function's If node subgraphs were version-converted
+        # Verify nodes inside the function's If node subgraphs were version-converted
         if_node = func[0]
         self.assertEqual(if_node.op_type, "If")
         self.assertEqual(if_node.version, 20)