onnx 1.15.0__cp311-cp311-win_amd64.whl → 1.16.1__cp311-cp311-win_amd64.whl

onnx/test/schema_test.py

@@ -1,8 +1,9 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+import contextlib
 import unittest
-from typing import Sequence
+from typing import List, Sequence
 
 import parameterized
 
@@ -258,5 +259,169 @@ class TestAttribute(unittest.TestCase):
         self.assertEqual("attr1 description", attribute.description)
 
 
+@parameterized.parameterized_class(
+    [
+        # register to exist domain
+        {
+            "op_type": "CustomOp",
+            "op_version": 5,
+            "op_domain": "",
+            "trap_op_version": [1, 2, 6, 7],
+        },
+        # register to new domain
+        {
+            "op_type": "CustomOp",
+            "op_version": 5,
+            "op_domain": "test",
+            "trap_op_version": [1, 2, 6, 7],
+        },
+    ]
+)
+class TestOpSchemaRegister(unittest.TestCase):
+    op_type: str
+    op_version: int
+    op_domain: str
+    # register some fake schema to check behavior
+    trap_op_version: List[int]
+
+    def setUp(self) -> None:
+        # Ensure the schema is unregistered
+        self.assertFalse(onnx.defs.has(self.op_type, self.op_domain))
+
+    def tearDown(self) -> None:
+        # Clean up the registered schema
+        for version in [*self.trap_op_version, self.op_version]:
+            with contextlib.suppress(onnx.defs.SchemaError):
+                onnx.defs.deregister_schema(self.op_type, version, self.op_domain)
+
+    def test_register_multi_schema(self):
+        for version in [*self.trap_op_version, self.op_version]:
+            op_schema = defs.OpSchema(
+                self.op_type,
+                self.op_domain,
+                version,
+            )
+            onnx.defs.register_schema(op_schema)
+            self.assertTrue(onnx.defs.has(self.op_type, version, self.op_domain))
+        for version in [*self.trap_op_version, self.op_version]:
+            # Also make sure the `op_schema` is accessible after register
+            registered_op = onnx.defs.get_schema(
+                op_schema.name, version, op_schema.domain
+            )
+            op_schema = defs.OpSchema(
+                self.op_type,
+                self.op_domain,
+                version,
+            )
+            self.assertEqual(str(registered_op), str(op_schema))
+
+    def test_using_the_specified_version_in_onnx_check(self):
+        input = f"""
+            <
+                ir_version: 7,
+                opset_import: [
+                    "{self.op_domain}" : {self.op_version}
+                ]
+            >
+            agraph (float[N, 128] X, int32 Y) => (float[N] Z)
+            {{
+                Z = {self.op_domain}.{self.op_type}<attr1=[1,2]>(X, Y)
+            }}
+           """
+        model = onnx.parser.parse_model(input)
+        op_schema = defs.OpSchema(
+            self.op_type,
+            self.op_domain,
+            self.op_version,
+            inputs=[
+                defs.OpSchema.FormalParameter("input1", "T"),
+                defs.OpSchema.FormalParameter("input2", "int32"),
+            ],
+            outputs=[
+                defs.OpSchema.FormalParameter("output1", "T"),
+            ],
+            type_constraints=[("T", ["tensor(float)"], "")],
+            attributes=[
+                defs.OpSchema.Attribute(
+                    "attr1", defs.OpSchema.AttrType.INTS, "attr1 description"
+                )
+            ],
+        )
+        with self.assertRaises(onnx.checker.ValidationError):
+            onnx.checker.check_model(model, check_custom_domain=True)
+        onnx.defs.register_schema(op_schema)
+        # The fake schema will raise check exception if selected in checker
+        for version in self.trap_op_version:
+            onnx.defs.register_schema(
+                defs.OpSchema(
+                    self.op_type,
+                    self.op_domain,
+                    version,
+                    outputs=[
+                        defs.OpSchema.FormalParameter("output1", "int32"),
+                    ],
+                )
+            )
+        onnx.checker.check_model(model, check_custom_domain=True)
+
+    def test_register_schema_raises_error_when_registering_a_schema_twice(self):
+        op_schema = defs.OpSchema(
+            self.op_type,
+            self.op_domain,
+            self.op_version,
+        )
+        onnx.defs.register_schema(op_schema)
+        with self.assertRaises(onnx.defs.SchemaError):
+            onnx.defs.register_schema(op_schema)
+
+    def test_deregister_the_specified_schema(self):
+        for version in [*self.trap_op_version, self.op_version]:
+            op_schema = defs.OpSchema(
+                self.op_type,
+                self.op_domain,
+                version,
+            )
+            onnx.defs.register_schema(op_schema)
+            self.assertTrue(onnx.defs.has(op_schema.name, version, op_schema.domain))
+        onnx.defs.deregister_schema(op_schema.name, self.op_version, op_schema.domain)
+        for version in self.trap_op_version:
+            self.assertTrue(onnx.defs.has(op_schema.name, version, op_schema.domain))
+        # Maybe has lesser op version in trap list
+        if onnx.defs.has(op_schema.name, self.op_version, op_schema.domain):
+            schema = onnx.defs.get_schema(
+                op_schema.name, self.op_version, op_schema.domain
+            )
+            self.assertLess(schema.since_version, self.op_version)
+
+    def test_deregister_schema_raises_error_when_opschema_does_not_exist(self):
+        with self.assertRaises(onnx.defs.SchemaError):
+            onnx.defs.deregister_schema(self.op_type, self.op_version, self.op_domain)
+
+    def test_legacy_schema_accessible_after_deregister(self):
+        op_schema = defs.OpSchema(
+            self.op_type,
+            self.op_domain,
+            self.op_version,
+        )
+        onnx.defs.register_schema(op_schema)
+        schema_a = onnx.defs.get_schema(
+            op_schema.name, op_schema.since_version, op_schema.domain
+        )
+        schema_b = onnx.defs.get_schema(op_schema.name, op_schema.domain)
+
+        def filter_schema(schemas):
+            return [op for op in schemas if op.name == op_schema.name]
+
+        schema_c = filter_schema(onnx.defs.get_all_schemas())
+        schema_d = filter_schema(onnx.defs.get_all_schemas_with_history())
+        self.assertEqual(len(schema_c), 1)
+        self.assertEqual(len(schema_d), 1)
+        # Avoid memory residue and access storage as much as possible
+        self.assertEqual(str(schema_a), str(op_schema))
+        self.assertEqual(str(schema_b), str(op_schema))
+        self.assertEqual(str(schema_c[0]), str(op_schema))
+        self.assertEqual(str(schema_d[0]), str(op_schema))
+
+
 if __name__ == "__main__":
     unittest.main()

onnx/test/shape_inference_test.py

@@ -395,6 +395,19 @@ class TestShapeInference(TestShapeInferenceHelper):
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
         )
 
+    @parameterized.expand(all_versions_for("Cast"))
+    @unittest.skip(
+        "Issue #5960"
+    )  # FIXME(#5960) propagateElemTypeFromAttributeToOutput does not validate against output type constraints
+    def test_cast_to_complex(self, _, version) -> None:
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (2, 4, 3))],
+            [make_node("Cast", ["x"], ["y"], to=TensorProto.COMPLEX128)],
+            [],
+        )
+
+        self.assertRaises(onnx.shape_inference.InferenceError, self._inferred, graph)
+
     @parameterized.expand(all_versions_for("CastLike"))
     def test_cast_like(self, _, version) -> None:
         graph = self._make_graph(
@@ -1652,6 +1665,21 @@ class TestShapeInference(TestShapeInferenceHelper):
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
         )
 
+    def test_squeeze_no_axes_opset11(self) -> None:
+        graph = self._make_graph(
+            [
+                ("x", TensorProto.FLOAT, (1, 3, 1, 1, 2, 1)),
+            ],
+            [make_node("Squeeze", ["x"], "y")],
+            [],
+        )
+        operatorsetid = OperatorSetIdProto()
+        operatorsetid.domain = ""
+        operatorsetid.version = 11
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 2))]
+        )
+
     def test_unsqueeze_regular(self) -> None:
         graph = self._make_graph(
             [("x", TensorProto.FLOAT, (3, 2)), ("axes", TensorProto.INT64, (4,))],
@@ -5250,11 +5278,14 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._make_matmulinteger_test((5, 1, 4, 2), (1, 3, 2, 3))
         self._make_matmulinteger_test((4, 2), (3, 2, 3))
 
-    def test_quantizelinear(self) -> None:
+    @parameterized.expand(
+        [onnx.TensorProto.FLOAT, onnx.TensorProto.FLOAT16, onnx.TensorProto.BFLOAT16]
+    )
+    def test_quantizelinear(self, elem_type) -> None:
         graph = self._make_graph(
             [
-                ("x", TensorProto.FLOAT, (30, 4, 5)),
-                ("y_scale", TensorProto.FLOAT, ()),
+                ("x", elem_type, (30, 4, 5)),
+                ("y_scale", elem_type, ()),
                 ("y_zero_point", TensorProto.UINT8, ()),
             ],
             [make_node("QuantizeLinear", ["x", "y_scale", "y_zero_point"], ["y"])],
@@ -5284,18 +5315,106 @@ class TestShapeInference(TestShapeInferenceHelper):
             graph, [make_tensor_value_info("y", TensorProto.UINT8, (30, 4, 5))]
         )
 
-    def test_dequantizelinear(self) -> None:
+    def test_quantizelinear_output_dtype(self) -> None:
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 4, 5)), ("y_scale", TensorProto.FLOAT, ())],
+            [
+                make_node(
+                    "QuantizeLinear",
+                    ["x", "y_scale"],
+                    ["y"],
+                    output_dtype=TensorProto.UINT4,
+                )
+            ],
+            [],
+        )
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.UINT4, (3, 4, 5))]
+        )
+
+    def test_quantizelinear_zp_output_dtype(self) -> None:
+        graph = self._make_graph(
+            [
+                ("x", TensorProto.FLOAT, (3, 4, 5)),
+                ("y_scale", TensorProto.FLOAT, ()),
+                ("y_zero_point", TensorProto.UINT16, ()),
+            ],
+            [
+                make_node(
+                    "QuantizeLinear",
+                    ["x", "y_scale", "y_zero_point"],
+                    ["y"],
+                    output_dtype=TensorProto.UINT16,
+                )
+            ],
+            [],
+        )
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.UINT16, (3, 4, 5))]
+        )
+
+    def test_quantizelinear_zp_output_dtype_conflicted(self) -> None:
+        graph = self._make_graph(
+            [
+                ("x", TensorProto.FLOAT, (3, 4, 5)),
+                ("y_scale", TensorProto.FLOAT, ()),
+                ("y_zero_point", TensorProto.UINT16, ()),
+            ],
+            [
+                make_node(
+                    "QuantizeLinear",
+                    ["x", "y_scale", "y_zero_point"],
+                    ["y"],
+                    output_dtype=TensorProto.INT4,
+                )
+            ],
+            [],
+        )
+
+        self.assertRaises(
+            onnx.shape_inference.InferenceError,
+            self._inferred,
+            graph,
+        )
+
+    @unittest.skip(
+        "Issue #5960"
+    )  # FIXME(#5960) propagateElemTypeFromAttributeToOutput does not validate against output type constraints
+    def test_quantizelinear_invalid_output_dtype(self) -> None:
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 4, 5)), ("y_scale", TensorProto.FLOAT, ())],
+            [
+                make_node(
+                    "QuantizeLinear",
+                    ["x", "y_scale"],
+                    ["y"],
+                    output_dtype=TensorProto.FLOAT16,
+                )
+            ],
+            [],
+        )
+
+        self.assertRaises(
+            onnx.shape_inference.InferenceError,
+            self._inferred,
+            graph,
+        )
+
+    @parameterized.expand(
+        [onnx.TensorProto.FLOAT, onnx.TensorProto.FLOAT16, onnx.TensorProto.BFLOAT16]
+    )
+    def test_dequantizelinear(self, elem_type) -> None:
         graph = self._make_graph(
             [
                 ("x", TensorProto.UINT8, (30, 4, 5)),
-                ("x_scale", TensorProto.FLOAT, ()),
+                ("x_scale", elem_type, ()),
                 ("x_zero_point", TensorProto.UINT8, ()),
             ],
             [make_node("DequantizeLinear", ["x", "x_scale", "x_zero_point"], ["y"])],
             [],
         )
         self._assert_inferred(
-            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (30, 4, 5))]
+            graph, [make_tensor_value_info("y", elem_type, (30, 4, 5))]
         )
 
     def test_dynamicquantizelinear(self) -> None:
@@ -9457,6 +9576,139 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
         )
 
+    @parameterized.expand([TensorProto.FLOAT, TensorProto.DOUBLE, TensorProto.FLOAT16])
+    @unittest.skipUnless(ONNX_ML, "ONNX_ML required to test ai.onnx.ml operators")
+    def test_tree_ensemble(self, dtype) -> None:
+        interior_nodes = 5
+        leaves = 9
+        tree = make_node(
+            "TreeEnsemble",
+            ["x"],
+            ["y"],
+            domain=ONNX_ML_DOMAIN,
+            n_targets=5,
+            nodes_featureids=[0] * interior_nodes,
+            nodes_splits=make_tensor(
+                "nodes_splits",
+                dtype,
+                (interior_nodes,),
+                list(range(interior_nodes)),
+            ),
+            nodes_modes=make_tensor(
+                "nodes_modes",
+                TensorProto.UINT8,
+                (interior_nodes,),
+                [0] * interior_nodes,
+            ),
+            nodes_truenodeids=[0] * interior_nodes,
+            nodes_falsenodeids=[0] * interior_nodes,
+            nodes_trueleafs=[0] * interior_nodes,
+            nodes_falseleafs=[0] * interior_nodes,
+            membership_values=make_tensor(
+                "membership_values",
+                dtype,
+                (7,),
+                [0.0, 0.1, 0.2, np.nan, 0.4, 0.5, 1.0],
+            ),
+            leaf_targetids=[0] * leaves,
+            leaf_weights=make_tensor("leaf_weights", dtype, (leaves,), [1] * leaves),
+            tree_roots=[0],
+        )
+
+        graph = self._make_graph(
+            [("x", dtype, ("Batch Size", "Features"))],
+            [tree],
+            [],
+        )
+
+        self._assert_inferred(
+            graph,
+            [make_tensor_value_info("y", dtype, ("Batch Size", 5))],
+            opset_imports=[
+                make_opsetid(ONNX_ML_DOMAIN, 5),
+                make_opsetid(ONNX_DOMAIN, 11),
+            ],
+        )
+
+    @parameterized.expand(
+        [
+            {
+                "nodes_truenodeids": [0] * 6,
+                "leaf_weights": make_tensor(
+                    "leaf_weights", TensorProto.DOUBLE, (9,), [1] * 9
+                ),
+                "nodes_splits": make_tensor(
+                    "nodes_splits", TensorProto.DOUBLE, (5,), [1] * 5
+                ),
+            },
+            {
+                "nodes_truenodeids": [0] * 5,
+                "leaf_weights": make_tensor(
+                    "leaf_weights", TensorProto.FLOAT, (9,), [1] * 9
+                ),
+                "nodes_splits": make_tensor(
+                    "nodes_splits", TensorProto.DOUBLE, (5,), [1] * 5
+                ),
+            },
+            {
+                "nodes_truenodeids": [0] * 5,
+                "leaf_weights": make_tensor(
+                    "leaf_weights", TensorProto.DOUBLE, (18,), [1] * 18
+                ),
+                "nodes_splits": make_tensor(
+                    "nodes_splits", TensorProto.DOUBLE, (5,), [1] * 5
+                ),
+            },
+            {
+                "nodes_truenodeids": [0] * 5,
+                "leaf_weights": make_tensor(
+                    "leaf_weights", TensorProto.DOUBLE, (9,), [1] * 9
+                ),
+                "nodes_splits": make_tensor(
+                    "nodes_splits", TensorProto.FLOAT, (5,), [1] * 5
+                ),
+            },
+        ]
+    )
+    @unittest.skipUnless(ONNX_ML, "ONNX_ML required to test ai.onnx.ml operators")
+    def test_tree_ensemble_fails_if_invalid_attributes(
+        self,
+        nodes_truenodeids,
+        leaf_weights,
+        nodes_splits,
+    ) -> None:
+        interior_nodes = 5
+        leaves = 9
+        tree = make_node(
+            "TreeEnsemble",
+            ["x"],
+            ["y"],
+            domain=ONNX_ML_DOMAIN,
+            n_targets=5,
+            nodes_featureids=[0] * interior_nodes,
+            nodes_splits=nodes_splits,
+            nodes_modes=make_tensor(
+                "nodes_modes",
+                TensorProto.UINT8,
+                (interior_nodes,),
+                [0] * interior_nodes,
+            ),
+            nodes_truenodeids=nodes_truenodeids,
+            nodes_falsenodeids=[0] * interior_nodes,
+            nodes_trueleafs=[0] * interior_nodes,
+            nodes_falseleafs=[0] * interior_nodes,
+            leaf_targetids=[0] * leaves,
+            leaf_weights=leaf_weights,
+            tree_roots=[0],
+        )
+
+        graph = self._make_graph(
+            [("x", TensorProto.DOUBLE, ("Batch Size", "Features"))],
+            [tree],
+            [],
+        )
+        self.assertRaises(onnx.shape_inference.InferenceError, self._inferred, graph)
+
     @unittest.skipUnless(ONNX_ML, "ONNX_ML required to test ai.onnx.ml operators")
     def test_tree_ensemble_classifier(self) -> None:
         tree = make_node(
@@ -9623,6 +9875,33 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(graph, [make_tensor_value_info("output", TensorProto.INT64, (2, "N", 3, None))])  # type: ignore
 
+    def test_check_type_when_schema_has_empty_io(self):
+        input = """
+            <
+                ir_version: 7,
+                opset_import: ["" : 1]
+            >
+            agraph (X, Y) => (Z)
+            {
+                Z = CustomOp(X, Y)
+            }
+           """
+        model = onnx.parser.parse_model(input)
+
+        op_schema = defs.OpSchema(
+            "CustomOp",
+            "",
+            1,
+            inputs=[],
+            outputs=[],
+        )
+        onnx.defs.register_schema(op_schema)
+        with self.assertRaises(onnx.shape_inference.InferenceError):
+            onnx.shape_inference.infer_shapes(model, True)
+        onnx.defs.deregister_schema(
+            op_schema.name, op_schema.since_version, op_schema.domain
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

onnx/test/symbolic_shape_test.py

@@ -14,8 +14,7 @@ class TestSymbolicShape(unittest.TestCase):
     def _assert_valueinfo_shape(
         self, onnx_model: ModelProto, value_infos: List[ValueInfoProto]
     ) -> None:
-        """
-        Assert onnx_model.value_info should be the same as expected value_infos
+        """Assert onnx_model.value_info should be the same as expected value_infos
         Instead of exact symbol, use -1 to represent symbolic shape in expected value_infos
         """
         for expected_vi in value_infos:
@@ -36,9 +35,7 @@ class TestSymbolicShape(unittest.TestCase):
                     assert dim.dim_value == expected_dim.dim_value, f"{onnx_model}"
 
     def _count_unique_dim_param_number(self, onnx_model: ModelProto) -> int:
-        """
-        return the total number of unique symbolic shape
-        """
+        """Return the total number of unique symbolic shape"""
         symbol_shape_set = set()
         inputs = list(onnx_model.graph.input)
         outputs = list(onnx_model.graph.output)
@@ -52,9 +49,7 @@ class TestSymbolicShape(unittest.TestCase):
     def _get_shape_from_name(
         self, onnx_model: ModelProto, name: str
     ) -> Optional[TensorShapeProto]:
-        """
-        Get shape from tensor_type or sparse_tensor_type according to given name
-        """
+        """Get shape from tensor_type or sparse_tensor_type according to given name"""
         inputs = list(onnx_model.graph.input)
         outputs = list(onnx_model.graph.output)
         valueinfos = list(onnx_model.graph.value_info)