onnx 1.17.0__cp311-cp311-win_amd64.whl → 1.18.0__cp311-cp311-win_amd64.whl

onnx/__init__.py

@@ -17,6 +17,7 @@ __all__ = [
     "IR_VERSION_2020_5_8",
     "IR_VERSION_2021_7_30",
     "IR_VERSION_2023_5_5",
+    "IR_VERSION_2024_3_25",
     "EXPERIMENTAL",
     "STABLE",
     # Modules
@@ -35,10 +36,13 @@ __all__ = [
     "version_converter",
     # Proto classes
     "AttributeProto",
+    "DeviceConfigurationProto",
     "FunctionProto",
     "GraphProto",
+    "IntIntListEntryProto",
     "MapProto",
     "ModelProto",
+    "NodeDeviceConfigurationProto",
     "NodeProto",
     "OperatorProto",
     "OperatorSetIdProto",
@@ -46,6 +50,9 @@ __all__ = [
     "OperatorStatus",
     "OptionalProto",
     "SequenceProto",
+    "SimpleShardedDimProto",
+    "ShardedDimProto",
+    "ShardingSpecProto",
     "SparseTensorProto",
     "StringStringEntryProto",
     "TensorAnnotation",
@@ -82,9 +89,11 @@ from onnx.external_data_helper import (
 )
 from onnx.onnx_pb import (
     AttributeProto,
+    DeviceConfigurationProto,
     EXPERIMENTAL,
     FunctionProto,
     GraphProto,
+    IntIntListEntryProto,
     IR_VERSION,
     IR_VERSION_2017_10_10,
     IR_VERSION_2017_10_30,
@@ -95,11 +104,16 @@ from onnx.onnx_pb import (
     IR_VERSION_2020_5_8,
     IR_VERSION_2021_7_30,
     IR_VERSION_2023_5_5,
+    IR_VERSION_2024_3_25,
     ModelProto,
+    NodeDeviceConfigurationProto,
     NodeProto,
     OperatorSetIdProto,
     OperatorStatus,
     STABLE,
+    SimpleShardedDimProto,
+    ShardedDimProto,
+    ShardingSpecProto,
     SparseTensorProto,
     StringStringEntryProto,
     TensorAnnotation,

onnx/_custom_element_types.py

@@ -52,6 +52,11 @@ uint4 = np.dtype((np.uint8, {"uint4": (np.uint8, 0)}))
 #: than its onnx size.
 int4 = np.dtype((np.int8, {"int4": (np.int8, 0)}))
 
+#: Defines float 4 e2m1 type, see See :ref:`onnx-detail-float4` for technical details.
+#: Do note that one integer is stored using a byte and therefore is twice bigger
+#: than its onnx size.
+float4e2m1 = np.dtype((np.uint8, {"float4e2m1": (np.uint8, 0)}))
+
 mapping_name_to_data_type = {
     "bfloat16": onnx.TensorProto.BFLOAT16,
     "e4m3fn": onnx.TensorProto.FLOAT8E4M3FN,
@@ -60,4 +65,5 @@ mapping_name_to_data_type = {
     "e5m2fnuz": onnx.TensorProto.FLOAT8E5M2FNUZ,
     "int4": onnx.TensorProto.INT4,
     "uint4": onnx.TensorProto.UINT4,
+    "float4e2m1": onnx.TensorProto.FLOAT4E2M1,
 }

onnx/backend/base.py

@@ -4,14 +4,17 @@
 from __future__ import annotations
 
 from collections import namedtuple
-from typing import Any, NewType, Sequence
-
-import numpy
+from typing import TYPE_CHECKING, Any, NewType
 
 import onnx.checker
 import onnx.onnx_cpp2py_export.checker as c_checker
 from onnx import IR_VERSION, ModelProto, NodeProto
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    import numpy
+
 
 class DeviceType:
     """Describes device type."""
@@ -76,14 +79,20 @@ class Backend:
 
     @classmethod
     def is_compatible(
-        cls, model: ModelProto, device: str = "CPU", **kwargs: Any  # noqa: ARG003
+        cls,
+        model: ModelProto,  # noqa: ARG003
+        device: str = "CPU",  # noqa: ARG003
+        **kwargs: Any,  # noqa: ARG003
     ) -> bool:
         # Return whether the model is compatible with the backend.
         return True
 
     @classmethod
     def prepare(
-        cls, model: ModelProto, device: str = "CPU", **kwargs: Any  # noqa: ARG003
+        cls,
+        model: ModelProto,
+        device: str = "CPU",  # noqa: ARG003
+        **kwargs: Any,  # noqa: ARG003
     ) -> BackendRep | None:
         # TODO Remove Optional from return type
         onnx.checker.check_model(model)

onnx/backend/sample/ops/__init__.py

@@ -7,7 +7,10 @@ import importlib
 import inspect
 import pkgutil
 import sys
-from types import ModuleType
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from types import ModuleType
 
 
 def collect_sample_implementations() -> dict[str, str]:
@@ -17,9 +20,9 @@ def collect_sample_implementations() -> dict[str, str]:
 
 
 def _recursive_scan(package: ModuleType, dict_: dict[str, str]) -> None:
-    pkg_dir = package.__path__  # type: ignore
+    pkg_dir = package.__path__
     module_location = package.__name__
-    for _module_loader, name, ispkg in pkgutil.iter_modules(pkg_dir):  # type: ignore
+    for _module_loader, name, ispkg in pkgutil.iter_modules(pkg_dir):
         module_name = f"{module_location}.{name}"  # Module/package
         module = importlib.import_module(module_name)
         dict_[name] = inspect.getsource(module)

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

@@ -4,14 +4,18 @@
 from __future__ import annotations
 
 import sys
-from typing import Sequence
+from typing import TYPE_CHECKING
 
-import numpy as np
-
-from onnx import ModelProto
 from onnx.backend.test.case.test_case import TestCase
 from onnx.backend.test.case.utils import import_recursive
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    import numpy as np
+
+    from onnx import ModelProto
+
 _SimpleModelTestCases = []
 
 

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

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from typing import Sequence
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -11,6 +11,9 @@ import onnx
 from onnx.backend.test.case.base import Base
 from onnx.backend.test.case.model import expect
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 
 class ExpandDynamicShape(Base):
     @staticmethod

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

@@ -66,10 +66,10 @@ def SplitToSequenceImpl(
     if np.isscalar(split):
         split_indices = [
             i * split + 1 for i in range(dim_size) if i * split + 1 < dim_size
-        ]  # type: ignore
+        ]
     else:
         split_indices = np.cumsum(split) + 1
-    return np.array_split(tensor, split_indices, axis)  # type: ignore
+    return np.array_split(tensor, split_indices, axis)
 
 
 def ConcatFromSequenceImpl(
@@ -159,11 +159,11 @@ class Sequence(Base):
                 seq_insert_node3,
                 seq_at_node,
             ],
-            [x_shape, y_shape, z_shape, [], []],  # type: ignore
-            [out_shape],  # type: ignore
+            [x_shape, y_shape, z_shape, [], []],
+            [out_shape],
             ["X", "Y", "Z", "pos", "pos_at"],
             ["out"],
-            [onnx.TensorProto.FLOAT] * 3 + [onnx.TensorProto.INT64] * 2,  # type: ignore
+            [onnx.TensorProto.FLOAT] * 3 + [onnx.TensorProto.INT64] * 2,
             [onnx.TensorProto.FLOAT],
             [pos, pos_at],
         )
@@ -206,11 +206,11 @@ class Sequence(Base):
 
         graph = make_graph(
             [seq_construct_node, seq_erase_node, seq_at_node],
-            [tensor_shape, tensor_shape, tensor_shape, [], []],  # type: ignore
-            [tensor_shape],  # type: ignore
+            [tensor_shape, tensor_shape, tensor_shape, [], []],
+            [tensor_shape],
             ["X", "Y", "Z", "pos_erase", "pos_at"],
             ["out"],
-            [onnx.TensorProto.FLOAT] * 3 + [onnx.TensorProto.INT64] * 2,  # type: ignore
+            [onnx.TensorProto.FLOAT] * 3 + [onnx.TensorProto.INT64] * 2,
             [onnx.TensorProto.FLOAT],
             [pos_erase, pos_at],
         )
@@ -261,11 +261,11 @@ class Sequence(Base):
 
         graph = make_graph(
             [seq_construct_node, seq_erase_node, seq_insert_node, seq_at_node],
-            [tensor_shape, tensor_shape, tensor_shape, [], [], []],  # type: ignore
-            [tensor_shape],  # type: ignore
+            [tensor_shape, tensor_shape, tensor_shape, [], [], []],
+            [tensor_shape],
             ["X", "Y", "Z", "pos_erase", "pos_insert", "pos_at"],
             ["out"],
-            [onnx.TensorProto.FLOAT] * 3 + [onnx.TensorProto.INT64] * 3,  # type: ignore
+            [onnx.TensorProto.FLOAT] * 3 + [onnx.TensorProto.INT64] * 3,
             [onnx.TensorProto.FLOAT],
             [pos_erase, pos_insert, pos_at],
         )
@@ -295,11 +295,11 @@ class Sequence(Base):
 
         graph = make_graph(
             [seq_construct_node, seq_concat_node],
-            [tensor_shape] * 3,  # type: ignore
-            [concat_out_shape],  # type: ignore
+            [tensor_shape] * 3,
+            [concat_out_shape],
             ["X", "Y", "Z"],
             ["out"],
-            [onnx.TensorProto.FLOAT] * 3,  # type: ignore
+            [onnx.TensorProto.FLOAT] * 3,
             [onnx.TensorProto.FLOAT],
         )
         model = onnx.helper.make_model_gen_version(
@@ -330,11 +330,11 @@ class Sequence(Base):
 
         graph = make_graph(
             [seq_construct_node, seq_concat_node],
-            [tensor_shape] * 3,  # type: ignore
-            [concat_out_shape],  # type: ignore
+            [tensor_shape] * 3,
+            [concat_out_shape],
             ["X", "Y", "Z"],
             ["out"],
-            [onnx.TensorProto.FLOAT] * 3,  # type: ignore
+            [onnx.TensorProto.FLOAT] * 3,
             [onnx.TensorProto.FLOAT],
         )
         model = onnx.helper.make_model_gen_version(
@@ -353,7 +353,7 @@ class Sequence(Base):
         seq_len_node = onnx.helper.make_node("SequenceLength", ["seq_1"], ["len"])
 
         tensor_shape = [2, 3, 4]
-        len_shape = []  # type: ignore
+        len_shape = []
 
         x = np.ones(tensor_shape, dtype=np.float32)
         out = SplitToSequenceImpl(x, axis=-1)
@@ -373,7 +373,7 @@ class Sequence(Base):
                     "len", onnx.TensorProto.INT64, len_shape
                 )
             ],
-        )  # type: ignore
+        )
 
         model = onnx.helper.make_model_gen_version(
             graph,
@@ -401,8 +401,8 @@ class Sequence(Base):
 
         graph = make_graph(
             [seq_split_node, seq_at_node],
-            [tensor_shape, []],  # type: ignore
-            [out_shape],  # type: ignore
+            [tensor_shape, []],
+            [out_shape],
             ["X", "pos_at"],
             ["out"],
             [onnx.TensorProto.DOUBLE, onnx.TensorProto.INT64],
@@ -422,8 +422,8 @@ class Sequence(Base):
         )
         seq_len_node = onnx.helper.make_node("SequenceLength", ["seq_1"], ["len"])
 
-        tensor_shape = ["n"]  # type: ignore
-        splits_shape = [3]  # type: ignore
+        tensor_shape = ["n"]
+        splits_shape = [3]
 
         x = np.array([]).astype(np.float32)
         splits = np.array([0, 0, 0]).astype(np.int64)
@@ -435,17 +435,17 @@ class Sequence(Base):
             inputs=[
                 onnx.helper.make_tensor_value_info(
                     "X", onnx.TensorProto.FLOAT, tensor_shape
-                ),  # type: ignore
+                ),
                 onnx.helper.make_tensor_value_info(
                     "Splits", onnx.TensorProto.INT64, splits_shape
                 ),
-            ],  # type: ignore
+            ],
             outputs=[
                 onnx.helper.make_tensor_value_info(
                     "len", onnx.TensorProto.INT64, len_shape
                 )
             ],
-        )  # type: ignore
+        )
 
         model = onnx.helper.make_model_gen_version(
             graph,

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

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from typing import Sequence
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -11,6 +11,9 @@ import onnx
 from onnx.backend.test.case.base import Base
 from onnx.backend.test.case.model import expect
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 
 class NormalizeStrings(Base):
     @staticmethod

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

@@ -7,9 +7,7 @@ import subprocess
 import sys
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, Callable, Sequence
-
-import numpy as np
+from typing import TYPE_CHECKING, Any, Callable
 
 import onnx
 from onnx.backend.test.case.test_case import TestCase
@@ -24,6 +22,11 @@ from onnx.onnx_pb import (
     TypeProto,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    import numpy as np
+
 _NodeTestCases = []
 _TargetOpType = None
 _DiffOpTypes = None
@@ -49,7 +52,7 @@ def _rename_edges_helper(
         if attr.HasField("ref_attr_name"):
             if attr.ref_attr_name in attribute_map:
                 new_attr = AttributeProto()
-                new_attr.CopyFrom(attribute_map[attr.ref_attr_name])  # type: ignore
+                new_attr.CopyFrom(attribute_map[attr.ref_attr_name])
                 new_attr.name = attr.name
                 new_node.attribute.extend([new_attr])
         else:
@@ -67,11 +70,11 @@ def _rename_edges_helper(
                 for sparse_init_desc in new_graph.sparse_initializer:
                     sg_rename[sparse_init_desc.values.name] = (
                         sparse_init_desc.values.name
-                    ) = (prefix + sparse_init_desc.values.name)
+                    ) = prefix + sparse_init_desc.values.name
                 for sparse_init_desc in new_graph.sparse_initializer:
                     sg_rename[sparse_init_desc.indices.name] = (
                         sparse_init_desc.indices.name
-                    ) = (prefix + sparse_init_desc.indices.name)
+                    ) = prefix + sparse_init_desc.indices.name
 
                 def subgraph_rename_helper(name: str) -> Any:
                     if name in sg_rename:  # noqa: B023
@@ -97,12 +100,10 @@ def function_expand_helper(
     io_names_map = {}
     attribute_map = {a.name: a for a in node.attribute}
 
-    for idx in range(len(function_proto.input)):
-        io_names_map[function_proto.input[idx]] = (
-            node.input[idx] if idx in range(len(node.input)) else ""
-        )
+    for idx, input in enumerate(function_proto.input):
+        io_names_map[input] = node.input[idx] if idx in range(len(node.input)) else ""
 
-    for idx in range(len(function_proto.output)):
+    for idx, output in enumerate(function_proto.output):
         # Even if the node has been created with optional outputs missing, we
         # can't assume that the function body handles this correctly, such as in
         # the case that output is also an intermediate value.
@@ -110,7 +111,7 @@ def function_expand_helper(
         # name will be generated if the missing output is used, the same as any
         # other internal tensor.
         if idx in range(len(node.output)) and node.output[idx] != "":
-            io_names_map[function_proto.output[idx]] = node.output[idx]
+            io_names_map[output] = node.output[idx]
 
     def rename_helper(internal_name: str) -> Any:
         if internal_name in io_names_map:
@@ -199,9 +200,7 @@ def _make_test_model_gen_version(graph: GraphProto, **kwargs: Any) -> ModelProto
         latest_onnx_version,
         latest_ml_version,
         latest_training_version,
-    ) = onnx.helper.VERSION_TABLE[-1][
-        2:5
-    ]  # type: ignore
+    ) = onnx.helper.VERSION_TABLE[-1][2:5]  # type: ignore
     if "opset_imports" in kwargs:
         for opset in kwargs["opset_imports"]:
             # If the test model uses an unreleased opset version (latest_version+1),
@@ -410,19 +409,19 @@ def get_diff_op_types():
         check=True,
     )
     # obtain list of added or modified files in this PR
-    obtain_diff = subprocess.Popen(
+    with subprocess.Popen(
         ["git", "diff", "--name-only", "--diff-filter=AM", "origin/main", "HEAD"],
         cwd=cwd_path,
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
-    )
-    stdoutput, _ = obtain_diff.communicate()
-    diff_list = stdoutput.split()
-    changed_op_types = []
-    for file in diff_list:
-        file_name = file.decode("utf-8")
-        if file_name.startswith("onnx/backend/test/case/node/") and file_name.endswith(
-            ".py"
-        ):
-            changed_op_types.append(file_name.split("/")[-1].replace(".py", ""))
-    return changed_op_types
+    ) as obtain_diff:
+        stdoutput, _ = obtain_diff.communicate()
+        diff_list = stdoutput.split()
+        changed_op_types = []
+        for file in diff_list:
+            file_name = file.decode("utf-8")
+            if file_name.startswith(
+                "onnx/backend/test/case/node/"
+            ) and file_name.endswith(".py"):
+                changed_op_types.append(file_name.split("/")[-1].replace(".py", ""))
+        return changed_op_types

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

@@ -6,7 +6,6 @@ from __future__ import annotations
 import numpy as np
 
 import onnx
-from onnx.backend.sample.ops.abs import abs
 from onnx.backend.test.case.base import Base
 from onnx.backend.test.case.node import expect
 
@@ -20,6 +19,6 @@ class Abs(Base):
             outputs=["y"],
         )
         x = np.random.randn(3, 4, 5).astype(np.float32)
-        y = abs(x)
+        y = np.abs(x)
 
         expect(node, inputs=[x], outputs=[y], name="test_abs")

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

@@ -11,7 +11,7 @@ from onnx.backend.test.case.node import expect
 from onnx.defs import AI_ONNX_PREVIEW_TRAINING_DOMAIN
 
 
-def apply_adagrad(r, t, x, g, h, norm_coefficient, epsilon, decay_factor):  # type: ignore
+def apply_adagrad(r, t, x, g, h, norm_coefficient, epsilon, decay_factor):
     # Compute adjusted learning-rate.
     r_ = r / (1 + t * decay_factor)
     # Add gradient of regularization term.

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

@@ -13,7 +13,7 @@ from onnx.defs import AI_ONNX_PREVIEW_TRAINING_DOMAIN
 
 def apply_adam(
     r, t, x, g, v, h, norm_coefficient, norm_coefficient_post, alpha, beta, epsilon
-):  # type: ignore
+):
     # Add gradient of regularization term.
     g_regularized = norm_coefficient * x + g
     # Update momentum.

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

@@ -23,18 +23,30 @@ class Add(Base):
         y = np.random.randn(3, 4, 5).astype(np.float32)
         expect(node, inputs=[x, y], outputs=[x + y], name="test_add")
 
-    @staticmethod
-    def export_add_uint8() -> None:
-        node = onnx.helper.make_node(
-            "Add",
-            inputs=["x", "y"],
-            outputs=["sum"],
-        )
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.int8)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.int8)
+        expect(node, inputs=[x, y], outputs=[x + y], name="test_add_int8")
+
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.int16)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.int16)
+        expect(node, inputs=[x, y], outputs=[x + y], name="test_add_int16")
 
         x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
         y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint8)
         expect(node, inputs=[x, y], outputs=[x + y], name="test_add_uint8")
 
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint16)
+        expect(node, inputs=[x, y], outputs=[x + y], name="test_add_uint16")
+
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint32)
+        expect(node, inputs=[x, y], outputs=[x + y], name="test_add_uint32")
+
+        x = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
+        y = np.random.randint(24, size=(3, 4, 5), dtype=np.uint64)
+        expect(node, inputs=[x, y], outputs=[x + y], name="test_add_uint64")
+
     @staticmethod
     def export_add_broadcast() -> None:
         node = onnx.helper.make_node(