[Core] introduce set_opset_import

Author: maltanar

src/qonnx/core/modelwrapper.py

@@ -759,3 +759,23 @@ def get_customop_wrapper(self, node, fallback_customop_version=util.get_preferre
                 f"using fallback_customop_version={fallback_customop_version}"
             )
             return getCustomOp(node, onnx_opset_version=fallback_customop_version)
+
+    def set_opset_import(self, domain, version):
+        """Sets the opset version for a given domain in the model's opset imports.
+        If the domain already exists, its version will be updated. If not, a new
+        opset import will be added.
+
+        Args:
+            domain (str): The domain name (e.g. "qonnx.custom_op.general")
+            version (int): The opset version number
+        """
+        # find if domain already exists in opset imports
+        for opset in self._model_proto.opset_import:
+            if opset.domain == domain:
+                opset.version = version
+                return
+        # domain not found, add new opset import
+        new_opset = onnx.OperatorSetIdProto()
+        new_opset.domain = domain
+        new_opset.version = version
+        self._model_proto.opset_import.append(new_opset)

tests/core/test_modelwrapper.py

@@ -33,7 +33,7 @@
 import qonnx.core.data_layout as DataLayout
 from qonnx.core.datatype import DataType
 from qonnx.core.modelwrapper import ModelWrapper
-from qonnx.util.basic import qonnx_make_model
+from qonnx.util.basic import get_preferred_onnx_opset, qonnx_make_model
 
 
 def test_modelwrapper():
@@ -231,3 +231,31 @@ def test_modelwrapper_set_tensor_shape_multiple_inputs():
     # check that order of inputs is preserved
     assert model.graph.input[0].name == "in1"
     assert model.graph.input[1].name == "in2"
+
+
+def test_modelwrapper_set_opset_import():
+    # Create a simple model
+    in1 = onnx.helper.make_tensor_value_info("in1", onnx.TensorProto.FLOAT, [4, 4])
+    out1 = onnx.helper.make_tensor_value_info("out1", onnx.TensorProto.FLOAT, [4, 4])
+    node = onnx.helper.make_node("Neg", inputs=["in1"], outputs=["out1"])
+    graph = onnx.helper.make_graph(
+        nodes=[node],
+        name="single_node_graph",
+        inputs=[in1],
+        outputs=[out1],
+    )
+    onnx_model = qonnx_make_model(graph, producer_name="opset-test-model")
+    model = ModelWrapper(onnx_model)
+
+    # Test setting new domain
+    model.set_opset_import("qonnx.custom_op.general", 1)
+    preferred_onnx_opset = get_preferred_onnx_opset()
+    assert model.get_opset_imports() == {"": preferred_onnx_opset, "qonnx.custom_op.general": 1}
+
+    # Test updating existing domain
+    model.set_opset_import("qonnx.custom_op.general", 2)
+    assert model.get_opset_imports() == {"": preferred_onnx_opset, "qonnx.custom_op.general": 2}
+
+    # Test setting ONNX main domain
+    model.set_opset_import("", 13)
+    assert model.get_opset_imports() == {"": 13, "qonnx.custom_op.general": 2}