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

onnx/__init__.py

@@ -16,6 +16,7 @@ __all__ = [
     "IR_VERSION_2019_9_19",
     "IR_VERSION_2020_5_8",
     "IR_VERSION_2021_7_30",
+    "IR_VERSION_2023_5_5",
     "EXPERIMENTAL",
     "STABLE",
     # Modules
@@ -93,6 +94,7 @@ from onnx.onnx_pb import (
     IR_VERSION_2019_9_19,
     IR_VERSION_2020_5_8,
     IR_VERSION_2021_7_30,
+    IR_VERSION_2023_5_5,
     ModelProto,
     NodeProto,
     OperatorSetIdProto,
@@ -186,7 +188,7 @@ def _get_serializer(
 
 def load_model(
     f: IO[bytes] | str | os.PathLike,
-    format: _SupportedFormat | None = None,
+    format: _SupportedFormat | None = None,  # noqa: A002
     load_external_data: bool = True,
 ) -> ModelProto:
     """Loads a serialized ModelProto into memory.
@@ -218,7 +220,7 @@ def load_model(
 
 def load_tensor(
     f: IO[bytes] | str | os.PathLike,
-    format: _SupportedFormat | None = None,
+    format: _SupportedFormat | None = None,  # noqa: A002
 ) -> TensorProto:
     """Loads a serialized TensorProto into memory.
 
@@ -237,7 +239,7 @@ def load_tensor(
 
 def load_model_from_string(
     s: bytes | str,
-    format: _SupportedFormat = _DEFAULT_FORMAT,
+    format: _SupportedFormat = _DEFAULT_FORMAT,  # noqa: A002
 ) -> ModelProto:
     """Loads a binary string (bytes) that contains serialized ModelProto.
 
@@ -256,7 +258,7 @@ def load_model_from_string(
 
 def load_tensor_from_string(
     s: bytes,
-    format: _SupportedFormat = _DEFAULT_FORMAT,
+    format: _SupportedFormat = _DEFAULT_FORMAT,  # noqa: A002
 ) -> TensorProto:
     """Loads a binary string (bytes) that contains serialized TensorProto.
 
@@ -276,7 +278,7 @@ def load_tensor_from_string(
 def save_model(
     proto: ModelProto | bytes,
     f: IO[bytes] | str | os.PathLike,
-    format: _SupportedFormat | None = None,
+    format: _SupportedFormat | None = None,  # noqa: A002
     *,
     save_as_external_data: bool = False,
     all_tensors_to_one_file: bool = True,
@@ -284,8 +286,7 @@ def save_model(
     size_threshold: int = 1024,
     convert_attribute: bool = False,
 ) -> None:
-    """
-    Saves the ModelProto to the specified path and optionally, serialize tensors with raw data as external data before saving.
+    """Saves the ModelProto to the specified path and optionally, serialize tensors with raw data as external data before saving.
 
     Args:
         proto: should be a in-memory ModelProto
@@ -330,10 +331,9 @@ def save_model(
 def save_tensor(
     proto: TensorProto,
     f: IO[bytes] | str | os.PathLike,
-    format: _SupportedFormat | None = None,
+    format: _SupportedFormat | None = None,  # noqa: A002
 ) -> None:
-    """
-    Saves the TensorProto to the specified path.
+    """Saves the TensorProto to the specified path.
 
     Args:
         proto: should be a in-memory TensorProto

onnx/backend/base.py

@@ -14,9 +14,7 @@ from onnx import IR_VERSION, ModelProto, NodeProto
 
 
 class DeviceType:
-    """
-    Describes device type.
-    """
+    """Describes device type."""
 
     _Type = NewType("_Type", int)
     CPU: _Type = _Type(0)
@@ -24,8 +22,7 @@ class DeviceType:
 
 
 class Device:
-    """
-    Describes device type and device id
+    """Describes device type and device id
     syntax: device_type:device_id(optional)
     example: 'CPU', 'CUDA', 'CUDA:1'
     """
@@ -56,8 +53,7 @@ def namedtupledict(
 
 
 class BackendRep:
-    """
-    BackendRep is the handle that a Backend returns after preparing to execute
+    """BackendRep is the handle that a Backend returns after preparing to execute
     a model repeatedly. Users will then pass inputs to the run function of
     BackendRep to retrieve the corresponding results.
     """
@@ -68,8 +64,7 @@ class BackendRep:
 
 
 class Backend:
-    """
-    Backend is the entity that will take an ONNX model with inputs,
+    """Backend is the entity that will take an ONNX model with inputs,
     perform a computation, and then return the output.
 
     For one-off execution, users can use run_node and run_model to obtain results quickly.
@@ -112,11 +107,16 @@ class Backend:
         **kwargs: Dict[str, Any],
     ) -> Optional[Tuple[Any, ...]]:
         """Simple run one operator and return the results.
+
         Args:
+            node: The node proto.
+            inputs: Inputs to the node.
+            device: The device to run on.
             outputs_info: a list of tuples, which contains the element type and
-            shape of each output. First element of the tuple is the dtype, and
-            the second element is the shape. More use case can be found in
-            https://github.com/onnx/onnx/blob/main/onnx/backend/test/runner/__init__.py
+                shape of each output. First element of the tuple is the dtype, and
+                the second element is the shape. More use case can be found in
+                https://github.com/onnx/onnx/blob/main/onnx/backend/test/runner/__init__.py
+            kwargs: Other keyword arguments.
         """
         # TODO Remove Optional from return type
         if "opset_version" in kwargs:
@@ -131,8 +131,7 @@ class Backend:
 
     @classmethod
     def supports_device(cls, device: str) -> bool:
-        """
-        Checks whether the backend is compiled with particular device support.
+        """Checks whether the backend is compiled with particular device support.
         In particular it's used in the testing suite.
         """
         return True

onnx/backend/sample/ops/abs.py

@@ -3,5 +3,5 @@
 import numpy as np
 
 
-def abs(input: np.ndarray) -> np.ndarray:
+def abs(input: np.ndarray) -> np.ndarray:  # noqa: A001
     return np.abs(input)  # type: ignore[no-any-return]

onnx/backend/test/case/model/__init__.py

@@ -42,7 +42,6 @@ BASE_URL = "onnx/backend/test/data/light/light_%s.onnx"
 
 def collect_testcases() -> List[TestCase]:
     """Collect model test cases defined in python/numpy code."""
-
     real_model_testcases = []
 
     model_tests = [

onnx/backend/test/case/node/ai_onnx_ml/tree_ensemble.py

@@ -0,0 +1,122 @@
+# Copyright (c) ONNX Project Contributors
+
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+
+import onnx
+from onnx.backend.test.case.base import Base
+from onnx.backend.test.case.node import expect
+from onnx.helper import make_tensor
+
+
+class TreeEnsemble(Base):
+    @staticmethod
+    def export_tree_ensemble_single_tree() -> None:
+        node = onnx.helper.make_node(
+            "TreeEnsemble",
+            ["X"],
+            ["Y"],
+            domain="ai.onnx.ml",
+            n_targets=2,
+            membership_values=None,
+            nodes_missing_value_tracks_true=None,
+            nodes_hitrates=None,
+            aggregate_function=1,
+            post_transform=0,
+            tree_roots=[0],
+            nodes_modes=make_tensor(
+                "nodes_modes",
+                onnx.TensorProto.UINT8,
+                (3,),
+                np.array([0, 0, 0], dtype=np.uint8),
+            ),
+            nodes_featureids=[0, 0, 0],
+            nodes_splits=make_tensor(
+                "nodes_splits",
+                onnx.TensorProto.DOUBLE,
+                (3,),
+                np.array([3.14, 1.2, 4.2], dtype=np.float64),
+            ),
+            nodes_truenodeids=[1, 0, 1],
+            nodes_trueleafs=[0, 1, 1],
+            nodes_falsenodeids=[2, 2, 3],
+            nodes_falseleafs=[0, 1, 1],
+            leaf_targetids=[0, 1, 0, 1],
+            leaf_weights=make_tensor(
+                "leaf_weights",
+                onnx.TensorProto.DOUBLE,
+                (4,),
+                np.array([5.23, 12.12, -12.23, 7.21], dtype=np.float64),
+            ),
+        )
+
+        x = np.array([1.2, 3.4, -0.12, 1.66, 4.14, 1.77], np.float64).reshape(3, 2)
+        y = np.array([[5.23, 0], [5.23, 0], [0, 12.12]], dtype=np.float64)
+        expect(
+            node,
+            inputs=[x],
+            outputs=[y],
+            name="test_ai_onnx_ml_tree_ensemble_single_tree",
+        )
+
+    @staticmethod
+    def export_tree_ensemble_set_membership() -> None:
+        node = onnx.helper.make_node(
+            "TreeEnsemble",
+            ["X"],
+            ["Y"],
+            domain="ai.onnx.ml",
+            n_targets=4,
+            aggregate_function=1,
+            membership_values=make_tensor(
+                "membership_values",
+                onnx.TensorProto.FLOAT,
+                (8,),
+                [1.2, 3.7, 8, 9, np.nan, 12, 7, np.nan],
+            ),
+            nodes_missing_value_tracks_true=None,
+            nodes_hitrates=None,
+            post_transform=0,
+            tree_roots=[0],
+            nodes_modes=make_tensor(
+                "nodes_modes",
+                onnx.TensorProto.UINT8,
+                (3,),
+                np.array([0, 6, 6], dtype=np.uint8),
+            ),
+            nodes_featureids=[0, 0, 0],
+            nodes_splits=make_tensor(
+                "nodes_splits",
+                onnx.TensorProto.FLOAT,
+                (3,),
+                np.array([11, 232344.0, np.nan], dtype=np.float32),
+            ),
+            nodes_trueleafs=[0, 1, 1],
+            nodes_truenodeids=[1, 0, 1],
+            nodes_falseleafs=[1, 0, 1],
+            nodes_falsenodeids=[2, 2, 3],
+            leaf_targetids=[0, 1, 2, 3],
+            leaf_weights=make_tensor(
+                "leaf_weights", onnx.TensorProto.FLOAT, (4,), [1, 10, 1000, 100]
+            ),
+        )
+
+        x = np.array([1.2, 3.4, -0.12, np.nan, 12, 7], np.float32).reshape(-1, 1)
+        expected = np.array(
+            [
+                [1, 0, 0, 0],
+                [0, 0, 0, 100],
+                [0, 0, 0, 100],
+                [0, 0, 1000, 0],
+                [0, 0, 1000, 0],
+                [0, 10, 0, 0],
+            ],
+            dtype=np.float32,
+        )
+        expect(
+            node,
+            inputs=[x],
+            outputs=[expected],
+            name="test_ai_onnx_ml_tree_ensemble_set_membership",
+        )

onnx/backend/test/case/node/averagepool.py

@@ -18,8 +18,7 @@ from onnx.reference.ops.op_pool_common import (
 class AveragePool(Base):
     @staticmethod
     def export_averagepool_2d_precomputed_pads() -> None:
-        """
-        input_shape: [1, 1, 5, 5]
+        """input_shape: [1, 1, 5, 5]
         output_shape: [1, 1, 5, 5]
         pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
         """
@@ -63,8 +62,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_2d_precomputed_pads_count_include_pad() -> None:
-        """
-        input_shape: [1, 1, 5, 5]
+        """input_shape: [1, 1, 5, 5]
         output_shape: [1, 1, 5, 5]
         pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
         """
@@ -112,8 +110,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_2d_precomputed_strides() -> None:
-        """
-        input_shape: [1, 1, 5, 5]
+        """input_shape: [1, 1, 5, 5]
         output_shape: [1, 1, 2, 2]
         """
         node = onnx.helper.make_node(
@@ -147,8 +144,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_2d_precomputed_same_upper() -> None:
-        """
-        input_shape: [1, 1, 5, 5]
+        """input_shape: [1, 1, 5, 5]
         output_shape: [1, 1, 3, 3]
         pad_shape: [2, 2] -> [1, 1, 1, 1] by axis
         """
@@ -186,8 +182,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_1d_default() -> None:
-        """
-        input_shape: [1, 3, 32]
+        """input_shape: [1, 3, 32]
         output_shape: [1, 3, 31]
         """
         node = onnx.helper.make_node(
@@ -211,8 +206,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_2d_default() -> None:
-        """
-        input_shape: [1, 3, 32, 32]
+        """input_shape: [1, 3, 32, 32]
         output_shape: [1, 3, 31, 31]
         """
         node = onnx.helper.make_node(
@@ -236,8 +230,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_3d_default() -> None:
-        """
-        input_shape: [1, 3, 32, 32, 32]
+        """input_shape: [1, 3, 32, 32, 32]
         output_shape: [1, 3, 31, 31, 31]
         """
         node = onnx.helper.make_node(
@@ -261,8 +254,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_2d_same_upper() -> None:
-        """
-        input_shape: [1, 3, 32, 32]
+        """input_shape: [1, 3, 32, 32]
         output_shape: [1, 3, 32, 32]
         pad_shape: [1, 1] -> [0, 1, 0, 1] by axis
         """
@@ -300,8 +292,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_2d_same_lower() -> None:
-        """
-        input_shape: [1, 3, 32, 32]
+        """input_shape: [1, 3, 32, 32]
         output_shape: [1, 3, 32, 32]
         pad_shape: [1, 1] -> [1, 0, 1, 0] by axis
         """
@@ -339,8 +330,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_2d_pads() -> None:
-        """
-        input_shape: [1, 3, 28, 28]
+        """input_shape: [1, 3, 28, 28]
         output_shape: [1, 3, 30, 30]
         pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
         """
@@ -375,8 +365,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_2d_pads_count_include_pad() -> None:
-        """
-        input_shape: [1, 3, 28, 28]
+        """input_shape: [1, 3, 28, 28]
         output_shape: [1, 3, 30, 30]
         pad_shape: [4, 4] -> [2, 2, 2, 2] by axis
         """
@@ -427,8 +416,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_2d_strides() -> None:
-        """
-        input_shape: [1, 3, 32, 32]
+        """input_shape: [1, 3, 32, 32]
         output_shape: [1, 3, 10, 10]
         """
         node = onnx.helper.make_node(
@@ -452,8 +440,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_2d_ceil() -> None:
-        """
-        input_shape: [1, 1, 4, 4]
+        """input_shape: [1, 1, 4, 4]
         output_shape: [1, 1, 2, 2]
         """
         node = onnx.helper.make_node(
@@ -482,8 +469,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_2d_dilations() -> None:
-        """
-        input_shape: [1, 1, 4, 4]
+        """input_shape: [1, 1, 4, 4]
         output_shape: [1, 1, 2, 2]
         """
         node = onnx.helper.make_node(
@@ -516,8 +502,7 @@ class AveragePool(Base):
 
     @staticmethod
     def export_averagepool_3d_dilations() -> None:
-        """
-        input_shape: [1, 1, 4, 4]
+        """input_shape: [1, 1, 4, 4]
         output_shape: [1, 1, 2, 2]
         """
         node = onnx.helper.make_node(

onnx/backend/test/case/node/cast.py

@@ -7,7 +7,8 @@ import sys
 import numpy as np
 
 import onnx
-from onnx import TensorProto, helper
+import onnx.reference.custom_element_types as custom
+from onnx import TensorProto, helper, subbyte
 from onnx.backend.test.case.base import Base
 from onnx.backend.test.case.node import expect
 from onnx.helper import (
@@ -50,10 +51,27 @@ class Cast(Base):
             ("FLOAT8E5M2", "FLOAT16"),
             ("FLOAT8E5M2FNUZ", "FLOAT"),
             ("FLOAT8E5M2FNUZ", "FLOAT16"),
+            ("FLOAT", "UINT4"),
+            ("FLOAT16", "UINT4"),
+            ("FLOAT", "INT4"),
+            ("FLOAT16", "INT4"),
+            ("UINT4", "FLOAT"),
+            ("UINT4", "FLOAT16"),
+            ("UINT4", "UINT8"),
+            ("INT4", "FLOAT"),
+            ("INT4", "FLOAT16"),
+            ("INT4", "INT8"),
         ]
 
         vect_float32_to_float8e4m3 = np.vectorize(float32_to_float8e4m3)
         vect_float32_to_float8e5m2 = np.vectorize(float32_to_float8e5m2)
+        vect_float32_to_uint4 = np.vectorize(
+            lambda x: subbyte.float32_to_4bit_unpacked(x, signed=False)
+        )
+        vect_float32_to_int4 = np.vectorize(
+            lambda x: subbyte.float32_to_4bit_unpacked(x, signed=True)
+        )
+
         f8_types = ("FLOAT8E4M3FN", "FLOAT8E4M3FNUZ", "FLOAT8E5M2", "FLOAT8E5M2FNUZ")
 
         for from_type, to_type in test_cases:
@@ -124,6 +142,9 @@ class Cast(Base):
                         "INF",
                         "+INF",
                         "-INF",
+                        "-0.0000001",
+                        "0.0000001",
+                        "-1000000",
                     ],
                     dtype=np.float32,
                 )
@@ -131,33 +152,33 @@ class Cast(Base):
                 if from_type == "FLOAT":
                     input_values = np_fp32
                     input = make_tensor(
-                        "x", TensorProto.FLOAT, [3, 4], np_fp32.tolist()
+                        "x", TensorProto.FLOAT, [3, 5], np_fp32.tolist()
                     )
                 elif from_type == "FLOAT16":
                     input_values = np_fp32.astype(np.float16).astype(np.float32)
                     input = make_tensor(
-                        "x", TensorProto.FLOAT16, [3, 4], input_values.tolist()
+                        "x", TensorProto.FLOAT16, [3, 5], input_values.tolist()
                     )
                 elif from_type == "FLOAT8E4M3FN":
                     input_values = float8e4m3_to_float32(
                         vect_float32_to_float8e4m3(np_fp32)
                     )
                     input = make_tensor(
-                        "x", TensorProto.FLOAT8E4M3FN, [3, 4], input_values.tolist()
+                        "x", TensorProto.FLOAT8E4M3FN, [3, 5], input_values.tolist()
                     )
                 elif from_type == "FLOAT8E4M3FNUZ":
                     input_values = float8e4m3_to_float32(
                         vect_float32_to_float8e4m3(np_fp32, uz=True), uz=True
                     )
                     input = make_tensor(
-                        "x", TensorProto.FLOAT8E4M3FNUZ, [3, 4], input_values.tolist()
+                        "x", TensorProto.FLOAT8E4M3FNUZ, [3, 5], input_values.tolist()
                     )
                 elif from_type == "FLOAT8E5M2":
                     input_values = float8e5m2_to_float32(
                         vect_float32_to_float8e5m2(np_fp32)
                     )
                     input = make_tensor(
-                        "x", TensorProto.FLOAT8E5M2, [3, 4], input_values.tolist()
+                        "x", TensorProto.FLOAT8E5M2, [3, 5], input_values.tolist()
                     )
                 elif from_type == "FLOAT8E5M2FNUZ":
                     input_values = float8e5m2_to_float32(
@@ -166,7 +187,7 @@ class Cast(Base):
                         uz=True,
                     )
                     input = make_tensor(
-                        "x", TensorProto.FLOAT8E5M2FNUZ, [3, 4], input_values.tolist()
+                        "x", TensorProto.FLOAT8E5M2FNUZ, [3, 5], input_values.tolist()
                     )
                 else:
                     raise ValueError(
@@ -200,9 +221,62 @@ class Cast(Base):
                         "Conversion from {from_type} to {to_type} is not tested."
                     )
                 expected_tensor = make_tensor(
-                    "x", getattr(TensorProto, to_type), [3, 4], expected.tolist()
+                    "x", getattr(TensorProto, to_type), [3, 5], expected.tolist()
+                )
+                output = expected_tensor
+            elif from_type in ("UINT4", "INT4") or to_type in ("UINT4", "INT4"):
+                np_fp32 = np.arange(-9, 16).astype(np.float32)
+                input_shape = (5, 5)
+                if from_type == "FLOAT":
+                    input_values = np_fp32
+                    input = make_tensor(
+                        "x", TensorProto.FLOAT, input_shape, input_values.tolist()
+                    )
+                elif from_type == "FLOAT16":
+                    input_values = np_fp32.astype(np.float16)
+                    input = make_tensor(
+                        "x", TensorProto.FLOAT16, input_shape, input_values.tolist()
+                    )
+                elif from_type == "UINT4":
+                    input_values = vect_float32_to_uint4(np_fp32)
+                    input = make_tensor(
+                        "x", TensorProto.UINT4, input_shape, input_values.tolist()
+                    )
+                elif from_type == "INT4":
+                    input_values = vect_float32_to_int4(np_fp32)
+                    input = make_tensor(
+                        "x", TensorProto.INT4, input_shape, input_values.tolist()
+                    )
+                else:
+                    raise ValueError(
+                        "Conversion from {from_type} to {to_type} is not tested."
+                    )
+                if to_type == "UINT4":
+                    expected = vect_float32_to_uint4(input_values).astype(custom.uint4)
+                elif to_type == "INT4":
+                    expected = vect_float32_to_int4(input_values).astype(custom.int4)
+                elif to_type == "FLOAT16":
+                    expected = input_values.astype(np.float16)
+                elif to_type == "FLOAT":
+                    expected = input_values
+                elif to_type == "UINT8":
+                    expected = input_values.astype(np.uint8)
+                elif to_type == "INT8":
+                    expected = input_values.astype(np.int8)
+                else:
+                    raise ValueError(
+                        "Conversion from {from_type} to {to_type} is not tested."
+                    )
+                expected_tensor = make_tensor(
+                    "y", getattr(TensorProto, to_type), input_shape, expected.tolist()
                 )
                 output = expected_tensor
+                input_type_proto = onnx.helper.make_tensor_type_proto(
+                    getattr(TensorProto, from_type), input_shape
+                )
+                output_type_proto = onnx.helper.make_tensor_type_proto(
+                    getattr(TensorProto, to_type), input_shape
+                )
 
             elif from_type != "STRING":
                 input = np.random.random_sample(shape).astype(
@@ -295,17 +369,20 @@ class Cast(Base):
                     "INF",
                     "+INF",
                     "-INF",
+                    "-0.0000001",
+                    "0.0000001",
+                    "-1000000",
                 ],
                 dtype=np.float32,
             )
 
             if from_type == "FLOAT":
                 input_values = np_fp32
-                input = make_tensor("x", TensorProto.FLOAT, [3, 4], np_fp32.tolist())
+                input = make_tensor("x", TensorProto.FLOAT, [3, 5], np_fp32.tolist())
             elif from_type == "FLOAT16":
                 input_values = np_fp32.astype(np.float16).astype(np.float32)
                 input = make_tensor(
-                    "x", TensorProto.FLOAT16, [3, 4], input_values.tolist()
+                    "x", TensorProto.FLOAT16, [3, 5], input_values.tolist()
                 )
             else:
                 raise ValueError(
@@ -333,7 +410,7 @@ class Cast(Base):
             tensor = TensorProto()
             tensor.data_type = getattr(TensorProto, to_type)
             tensor.name = "x"
-            tensor.dims.extend([3, 4])
+            tensor.dims.extend([3, 5])
             field = tensor_dtype_to_field(tensor.data_type)
             getattr(tensor, field).extend(ivals)