onnx 1.16.2__cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl → 1.18.0__cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.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,
@@ -134,7 +148,9 @@ from onnx import (
 # Supported model formats that can be loaded from and saved to
 # The literals are formats with built-in support. But we also allow users to
 # register their own formats. So we allow str as well.
-_SupportedFormat = Union[Literal["protobuf", "textproto"], str]
+_SupportedFormat = Union[
+    Literal["protobuf", "textproto", "onnxtxt", "json"], str  # noqa: PYI051
+]
 # Default serialization format
 _DEFAULT_FORMAT = "protobuf"
 

onnx/_custom_element_types.py

@@ -0,0 +1,69 @@
+# Copyright (c) ONNX Project Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+"""This module defines custom dtypes not supported by numpy.
+Function :func:`onnx.numpy_helper.from_array`
+and :func:`onnx.numpy_helper.to_array` are using them
+to convert arrays from/to these types.
+Class :class:`onnx.reference.ReferenceEvalutor` also uses them.
+To create such an array for unit test for example, it is convenient to write
+something like the following:
+
+.. exec_code::
+
+    import numpy as np
+    from onnx import TensorProto
+    from onnx.reference.ops.op_cast import Cast_19 as Cast
+
+    tensor_bfloat16 = Cast.eval(np.array([0, 1], dtype=np.float32), to=TensorProto.BFLOAT16)
+
+The numpy representation dtypes used below are meant for internal use. They may change in the
+future based on the industry standardization of these numpy types.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+import onnx
+
+#: Defines a bfloat16 as a uint16.
+bfloat16 = np.dtype((np.uint16, {"bfloat16": (np.uint16, 0)}))
+
+#: Defines float 8 e4m3fn type, see See :ref:`onnx-detail-float8` for technical details.
+float8e4m3fn = np.dtype((np.uint8, {"e4m3fn": (np.uint8, 0)}))
+
+#: Defines float 8 e4m3fnuz type, see See :ref:`onnx-detail-float8` for technical details.
+float8e4m3fnuz = np.dtype((np.uint8, {"e4m3fnuz": (np.uint8, 0)}))
+
+#: Defines float 8 e5m2 type, see See :ref:`onnx-detail-float8` for technical details.
+float8e5m2 = np.dtype((np.uint8, {"e5m2": (np.uint8, 0)}))
+
+#: Defines float 8 e5m2fnuz type, see See :ref:`onnx-detail-float8` for technical details.
+float8e5m2fnuz = np.dtype((np.uint8, {"e5m2fnuz": (np.uint8, 0)}))
+
+#: Defines int4, see See :ref:`onnx-detail-int4` for technical details.
+#: Do note that one integer is stored using a byte and therefore is twice bigger
+#: than its onnx size.
+uint4 = np.dtype((np.uint8, {"uint4": (np.uint8, 0)}))
+
+#: Defines int4, see See :ref:`onnx-detail-int4` for technical details.
+#: Do note that one integer is stored using a byte and therefore is twice bigger
+#: 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,
+    "e4m3fnuz": onnx.TensorProto.FLOAT8E4M3FNUZ,
+    "e5m2": onnx.TensorProto.FLOAT8E5M2,
+    "e5m2fnuz": onnx.TensorProto.FLOAT8E5M2FNUZ,
+    "int4": onnx.TensorProto.INT4,
+    "uint4": onnx.TensorProto.UINT4,
+    "float4e2m1": onnx.TensorProto.FLOAT4E2M1,
+}

onnx/backend/base.py

@@ -1,17 +1,20 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 from collections import namedtuple
-from typing import Any, Dict, NewType, Optional, Sequence, Tuple, Type
-
-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."""
@@ -37,11 +40,11 @@ class Device:
 
 def namedtupledict(
     typename: str, field_names: Sequence[str], *args: Any, **kwargs: Any
-) -> Type[Tuple[Any, ...]]:
+) -> type[tuple[Any, ...]]:
     field_names_map = {n: i for i, n in enumerate(field_names)}
     # Some output names are invalid python identifier, e.g. "0"
     kwargs.setdefault("rename", True)
-    data = namedtuple(typename, field_names, *args, **kwargs)  # type: ignore
+    data = namedtuple(typename, field_names, *args, **kwargs)  # type: ignore  # noqa: PYI024
 
     def getitem(self: Any, key: Any) -> Any:
         if isinstance(key, str):
@@ -58,7 +61,7 @@ class BackendRep:
     BackendRep to retrieve the corresponding results.
     """
 
-    def run(self, inputs: Any, **kwargs: Any) -> Tuple[Any, ...]:
+    def run(self, inputs: Any, **kwargs: Any) -> tuple[Any, ...]:  # noqa: ARG002
         """Abstract function."""
         return (None,)
 
@@ -76,15 +79,21 @@ class Backend:
 
     @classmethod
     def is_compatible(
-        cls, model: ModelProto, device: str = "CPU", **kwargs: Any
+        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
-    ) -> Optional[BackendRep]:
+        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)
         return None
@@ -92,7 +101,7 @@ class Backend:
     @classmethod
     def run_model(
         cls, model: ModelProto, inputs: Any, device: str = "CPU", **kwargs: Any
-    ) -> Tuple[Any, ...]:
+    ) -> tuple[Any, ...]:
         backend = cls.prepare(model, device, **kwargs)
         assert backend is not None
         return backend.run(inputs)
@@ -101,11 +110,13 @@ class Backend:
     def run_node(
         cls,
         node: NodeProto,
-        inputs: Any,
-        device: str = "CPU",
-        outputs_info: Optional[Sequence[Tuple[numpy.dtype, Tuple[int, ...]]]] = None,
-        **kwargs: Dict[str, Any],
-    ) -> Optional[Tuple[Any, ...]]:
+        inputs: Any,  # noqa: ARG003
+        device: str = "CPU",  # noqa: ARG003
+        outputs_info: (  # noqa: ARG003
+            Sequence[tuple[numpy.dtype, tuple[int, ...]]] | None
+        ) = None,
+        **kwargs: dict[str, Any],
+    ) -> tuple[Any, ...] | None:
         """Simple run one operator and return the results.
 
         Args:
@@ -130,7 +141,7 @@ class Backend:
         return None
 
     @classmethod
-    def supports_device(cls, device: str) -> bool:
+    def supports_device(cls, device: str) -> bool:  # noqa: ARG003
         """Checks whether the backend is compiled with particular device support.
         In particular it's used in the testing suite.
         """

onnx/backend/sample/ops/__init__.py

@@ -1,25 +1,28 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import importlib
 import inspect
 import pkgutil
 import sys
-from types import ModuleType
-from typing import Dict
+from typing import TYPE_CHECKING
 
+if TYPE_CHECKING:
+    from types import ModuleType
 
-def collect_sample_implementations() -> Dict[str, str]:
-    dict_: Dict[str, str] = {}
+
+def collect_sample_implementations() -> dict[str, str]:
+    dict_: dict[str, str] = {}
     _recursive_scan(sys.modules[__name__], dict_)
     return dict_
 
 
-def _recursive_scan(package: ModuleType, dict_: Dict[str, str]) -> None:
-    pkg_dir = package.__path__  # type: ignore
+def _recursive_scan(package: ModuleType, dict_: dict[str, str]) -> None:
+    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/sample/ops/abs.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import numpy as np
 

onnx/backend/test/__init__.py

@@ -1,6 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 __all__ = ["BackendTest"]
 # for backward compatibility

onnx/backend/test/case/__init__.py

@@ -1,14 +1,14 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import sys
-from typing import Dict, List, Tuple
 
 from onnx.backend.test.case.base import Snippets
 from onnx.backend.test.case.utils import import_recursive
 
 
-def collect_snippets() -> Dict[str, List[Tuple[str, str]]]:
+def collect_snippets() -> dict[str, list[tuple[str, str]]]:
     import_recursive(sys.modules[__name__])
     return Snippets

onnx/backend/test/case/base.py

@@ -1,16 +1,17 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import inspect
 from collections import defaultdict
 from textwrap import dedent
-from typing import Any, ClassVar, Dict, List, Tuple, Type
+from typing import Any, ClassVar
 
 import numpy as np
 
 
-def process_snippet(op_name: str, name: str, export: Any) -> Tuple[str, str]:
+def process_snippet(op_name: str, name: str, export: Any) -> tuple[str, str]:
     snippet_name = name[len("export_") :] or op_name.lower()
     source_code = dedent(inspect.getsource(export))
     # remove the function signature line
@@ -20,14 +21,14 @@ def process_snippet(op_name: str, name: str, export: Any) -> Tuple[str, str]:
     return snippet_name, dedent("\n".join(lines[2:]))
 
 
-Snippets: Dict[str, List[Tuple[str, str]]] = defaultdict(list)
+Snippets: dict[str, list[tuple[str, str]]] = defaultdict(list)
 
 
 class _Exporter(type):
-    exports: ClassVar[Dict[str, List[Tuple[str, str]]]] = defaultdict(list)
+    exports: ClassVar[dict[str, list[tuple[str, str]]]] = defaultdict(list)
 
     def __init__(
-        cls, name: str, bases: Tuple[Type[Any], ...], dct: Dict[str, Any]
+        cls, name: str, bases: tuple[type[Any], ...], dct: dict[str, Any]
     ) -> None:
         for k, v in dct.items():
             if k.startswith("export"):

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

@@ -1,16 +1,21 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import sys
-from typing import List, Optional, 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 = []
 
 
@@ -18,7 +23,7 @@ def expect(
     model: ModelProto,
     inputs: Sequence[np.ndarray],
     outputs: Sequence[np.ndarray],
-    name: Optional[str] = None,
+    name: str | None = None,
 ) -> None:
     name = name or model.graph.name
     _SimpleModelTestCases.append(
@@ -40,7 +45,7 @@ def expect(
 BASE_URL = "onnx/backend/test/data/light/light_%s.onnx"
 
 
-def collect_testcases() -> List[TestCase]:
+def collect_testcases() -> list[TestCase]:
     """Collect model test cases defined in python/numpy code."""
     real_model_testcases = []
 

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

@@ -1,8 +1,9 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
-from typing import Sequence
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -10,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/gradient.py

@@ -1,6 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import numpy as np
 

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

@@ -64,10 +64,12 @@ def SplitToSequenceImpl(
             results = np.array_split(tensor, split_indices, axis)
             return [np.squeeze(res, axis) for res in results]
     if np.isscalar(split):
-        split_indices = [i * split + 1 for i in range(dim_size) if i * split + 1 < dim_size]  # type: ignore
+        split_indices = [
+            i * split + 1 for i in range(dim_size) if i * split + 1 < dim_size
+        ]
     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(
@@ -157,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],
         )
@@ -204,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],
         )
@@ -259,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],
         )
@@ -293,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(
@@ -328,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(
@@ -351,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)
@@ -371,7 +373,7 @@ class Sequence(Base):
                     "len", onnx.TensorProto.INT64, len_shape
                 )
             ],
-        )  # type: ignore
+        )
 
         model = onnx.helper.make_model_gen_version(
             graph,
@@ -399,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],
@@ -420,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)
@@ -433,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/shrink.py

@@ -1,6 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import numpy as np
 

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

@@ -1,6 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import numpy as np
 

onnx/backend/test/case/model/single-relu.py

@@ -1,6 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import numpy as np
 

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

@@ -1,9 +1,9 @@
 # Copyright (c) ONNX Project Contributors
 
 # 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