Better support for ir_version (#82)

xadupre · web-flow · commit a54de21f2ea6 · 2024-03-26T13:30:08.000+01:00
* fixes for ir_version

* fix ut

* fix ut
diff --git a/_unittests/ut_light_api/test_backend_export.py b/_unittests/ut_light_api/test_backend_export.py
@@ -242,7 +242,7 @@ def run_node(cls, node, inputs, device=None, outputs_info=None, **kwargs):
 
 # The following tests are too slow with the reference implementation (Conv).
 backend_test.exclude(
-    "(FLOAT8|BFLOAT16|_opt_|_3d_|_momentum_|_4d_"
+    "(FLOAT8|BFLOAT16|INT4|_opt_|_3d_|_momentum_|_4d_|int4"
     "|test_adagrad"
     "|test_adam"
     "|test_ai_onnx_ml_"
@@ -270,6 +270,8 @@ def run_node(cls, node, inputs, device=None, outputs_info=None, **kwargs):
     "|test_squeezenet"
     "|test_vgg19"
     "|test_zfnet512"
+    "|test_range_float_type_positive_delta_expanded"
+    "|test_range_int32_type_negative_delta_expanded"
     ")"
 )
 
diff --git a/_unittests/ut_reference/test_backend_extended_reference_evaluator.py b/_unittests/ut_reference/test_backend_extended_reference_evaluator.py
@@ -149,7 +149,7 @@ def run_node(cls, node, inputs, device=None, outputs_info=None, **kwargs):
     "|test_scan_sum)"
 )
 
-if onnx_opset_version() < 21:
+if onnx_opset_version() < 200:
     # The following tests are using types not supported by NumPy.
     # They could be if method to_array is extended to support custom
     # types the same as the reference implementation does
@@ -164,8 +164,10 @@ def run_node(cls, node, inputs, device=None, outputs_info=None, **kwargs):
         "|test_cast_no_saturate_"
         "|_to_FLOAT8"
         "|_FLOAT8"
+        "|INT4"
         "|test_quantizelinear_e4m3fn"
         "|test_quantizelinear_e5m2"
+        "|test_scatter_with"
         ")"
     )
 
diff --git a/onnx_array_api/graph_api/graph_builder.py b/onnx_array_api/graph_api/graph_builder.py
@@ -156,6 +156,7 @@ def __init__(
         optimization_options: Optional[OptimizationOptions] = None,
         args: Optional[List[Any]] = None,
         verbose: int = 0,
+        ir_version: Optional[int] = None,
     ):
         self.optimization_options = optimization_options or OptimizationOptions()
         self.as_function = as_function
@@ -170,6 +171,7 @@ def __init__(
                 if isinstance(target_opset_or_existing_proto, int)
                 else target_opset_or_existing_proto
             )
+            self.ir_version = ir_version
             self.nodes = []
             self.initializers_dict = {}
             self.inputs = []
@@ -186,6 +188,7 @@ def __init__(
             ), "input_names must be empty if the input is an existing model."
             proto = target_opset_or_existing_proto
             self.opsets = {d.domain: d.version for d in proto.opset_import}
+            self.ir_version = ir_version or target_opset_or_existing_proto.ir_version
             self.nodes = list(proto.graph.node)
             self.initializers_dict = {i.name: i for i in proto.graph.initializer}
             self.initializers_dict.update(
@@ -674,6 +677,8 @@ def to_onnx(
         if self.verbose:
             print("[GraphBuilder] onh.make_model")
         model = oh.make_model(graph, opset_imports=opsets)
+        if self.ir_version:
+            model.ir_version = self.ir_version
         return model
 
     def _check_order_node(self, ind: int, node: NodeProto, existing: Set[str]):
