onnx 1.18.0__cp311-cp311-win_arm64.whl → 1.19.1rc1__cp311-cp311-win_arm64.whl

onnx/__init__.py

@@ -27,7 +27,6 @@ __all__ = [
     "gen_proto",
     "helper",
     "hub",
-    "mapping",
     "numpy_helper",
     "parser",
     "printer",
@@ -126,7 +125,7 @@ from onnx.onnx_pb import (
 )
 from onnx.onnx_operators_pb import OperatorProto, OperatorSetProto
 from onnx.onnx_data_pb import MapProto, OptionalProto, SequenceProto
-from onnx.version import version as __version__
+import onnx.version
 
 # Import common subpackages so they're available when you 'import onnx'
 from onnx import (
@@ -136,7 +135,6 @@ from onnx import (
     gen_proto,
     helper,
     hub,
-    mapping,
     numpy_helper,
     parser,
     printer,
@@ -145,6 +143,11 @@ from onnx import (
     version_converter,
 )
 
+if typing.TYPE_CHECKING:
+    from collections.abc import Sequence
+
+__version__ = onnx.version.version
+
 # 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.
@@ -156,20 +159,20 @@ _DEFAULT_FORMAT = "protobuf"
 
 
 def _load_bytes(f: IO[bytes] | str | os.PathLike) -> bytes:
-    if hasattr(f, "read") and callable(typing.cast(IO[bytes], f).read):
-        content = typing.cast(IO[bytes], f).read()
+    if hasattr(f, "read") and callable(typing.cast("IO[bytes]", f).read):
+        content = typing.cast("IO[bytes]", f).read()
     else:
-        f = typing.cast(Union[str, os.PathLike], f)
+        f = typing.cast("Union[str, os.PathLike]", f)
         with open(f, "rb") as readable:
             content = readable.read()
     return content
 
 
 def _save_bytes(content: bytes, f: IO[bytes] | str | os.PathLike) -> None:
-    if hasattr(f, "write") and callable(typing.cast(IO[bytes], f).write):
-        typing.cast(IO[bytes], f).write(content)
+    if hasattr(f, "write") and callable(typing.cast("IO[bytes]", f).write):
+        typing.cast("IO[bytes]", f).write(content)
     else:
-        f = typing.cast(Union[str, os.PathLike], f)
+        f = typing.cast("Union[str, os.PathLike]", f)
         with open(f, "wb") as writable:
             writable.write(content)
 
@@ -368,3 +371,98 @@ def save_tensor(
 load = load_model
 load_from_string = load_model_from_string
 save = save_model
+
+
+def _model_proto_repr(self: ModelProto) -> str:
+    if self.domain:
+        domain = f", domain='{self.domain}'"
+    else:
+        domain = ""
+    if self.producer_name:
+        producer_name = f", producer_name='{self.producer_name}'"
+    else:
+        producer_name = ""
+    if self.producer_version:
+        producer_version = f", producer_version='{self.producer_version}'"
+    else:
+        producer_version = ""
+    if self.graph:
+        graph = f", graph={self.graph!r}"
+    else:
+        graph = ""
+    if self.functions:
+        functions = f", functions=<{len(self.functions)} functions>"
+    else:
+        functions = ""
+    if self.opset_import:
+        opset_import = f", opset_import={_operator_set_protos_repr(self.opset_import)}"
+    else:
+        opset_import = ""
+    return f"ModelProto(ir_version={self.ir_version}{opset_import}{domain}{producer_name}{producer_version}{graph}{functions})"
+
+
+def _graph_proto_repr(self: GraphProto) -> str:
+    if self.initializer:
+        initializer = f", initializer=<{len(self.initializer)} initializers>"
+    else:
+        initializer = ""
+    if self.node:
+        node = f", node=<{len(self.node)} nodes>"
+    else:
+        node = ""
+    if self.value_info:
+        value_info = f", value_info=<{len(self.value_info)} value_info>"
+    else:
+        value_info = ""
+    if self.input:
+        input = f", input=<{len(self.input)} inputs>"
+    else:
+        input = ""
+    if self.output:
+        output = f", output=<{len(self.output)} outputs>"
+    else:
+        output = ""
+    return f"GraphProto('{self.name}'{input}{output}{initializer}{node}{value_info})"
+
+
+def _function_proto_repr(self: FunctionProto) -> str:
+    if self.domain:
+        domain = f", domain='{self.domain}'"
+    else:
+        domain = ""
+    if self.overload:
+        overload = f", overload='{self.overload}'"
+    else:
+        overload = ""
+    if self.node:
+        node = f", node=<{len(self.node)} nodes>"
+    else:
+        node = ""
+    if self.attribute:
+        attribute = f", attribute={self.attribute}"
+    else:
+        attribute = ""
+    if self.opset_import:
+        opset_import = f", opset_import={_operator_set_protos_repr(self.opset_import)}"
+    else:
+        opset_import = ""
+    if self.input:
+        input = f", input=<{len(self.input)} inputs>"
+    else:
+        input = ""
+    if self.output:
+        output = f", output=<{len(self.output)} outputs>"
+    else:
+        output = ""
+    return f"FunctionProto('{self.name}'{domain}{overload}{opset_import}{input}{output}{attribute}{node})"
+
+
+def _operator_set_protos_repr(protos: Sequence[OperatorSetIdProto]) -> str:
+    opset_imports = {proto.domain: proto.version for proto in protos}
+    return repr(opset_imports)
+
+
+# Override __repr__ for some proto classes to make it more efficient
+ModelProto.__repr__ = _model_proto_repr  # type: ignore[method-assign,assignment]
+GraphProto.__repr__ = _graph_proto_repr  # type: ignore[method-assign,assignment]
+FunctionProto.__repr__ = _function_proto_repr  # type: ignore[method-assign,assignment]

onnx/_mapping.py

@@ -0,0 +1,113 @@
+# Copyright (c) ONNX Project Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import ml_dtypes
+import numpy as np
+
+from onnx.onnx_pb import TensorProto
+
+
+class TensorDtypeMap(NamedTuple):
+    np_dtype: np.dtype
+    storage_dtype: int
+    name: str
+
+
+# tensor_dtype: (numpy type, storage type, string name)
+# The storage type is the type used to store the tensor in the *_data field of
+# a TensorProto. All available fields are float_data, int32_data, int64_data,
+# string_data, uint64_data and double_data.
+TENSOR_TYPE_MAP: dict[int, TensorDtypeMap] = {
+    int(TensorProto.FLOAT): TensorDtypeMap(
+        np.dtype("float32"), int(TensorProto.FLOAT), "TensorProto.FLOAT"
+    ),
+    int(TensorProto.UINT8): TensorDtypeMap(
+        np.dtype("uint8"), int(TensorProto.INT32), "TensorProto.UINT8"
+    ),
+    int(TensorProto.INT8): TensorDtypeMap(
+        np.dtype("int8"), int(TensorProto.INT32), "TensorProto.INT8"
+    ),
+    int(TensorProto.UINT16): TensorDtypeMap(
+        np.dtype("uint16"), int(TensorProto.INT32), "TensorProto.UINT16"
+    ),
+    int(TensorProto.INT16): TensorDtypeMap(
+        np.dtype("int16"), int(TensorProto.INT32), "TensorProto.INT16"
+    ),
+    int(TensorProto.INT32): TensorDtypeMap(
+        np.dtype("int32"), int(TensorProto.INT32), "TensorProto.INT32"
+    ),
+    int(TensorProto.INT64): TensorDtypeMap(
+        np.dtype("int64"), int(TensorProto.INT64), "TensorProto.INT64"
+    ),
+    int(TensorProto.BOOL): TensorDtypeMap(
+        np.dtype("bool"), int(TensorProto.INT32), "TensorProto.BOOL"
+    ),
+    int(TensorProto.FLOAT16): TensorDtypeMap(
+        np.dtype("float16"), int(TensorProto.INT32), "TensorProto.FLOAT16"
+    ),
+    int(TensorProto.BFLOAT16): TensorDtypeMap(
+        np.dtype(ml_dtypes.bfloat16),
+        int(TensorProto.INT32),
+        "TensorProto.BFLOAT16",
+    ),
+    int(TensorProto.DOUBLE): TensorDtypeMap(
+        np.dtype("float64"), int(TensorProto.DOUBLE), "TensorProto.DOUBLE"
+    ),
+    int(TensorProto.COMPLEX64): TensorDtypeMap(
+        np.dtype("complex64"), int(TensorProto.FLOAT), "TensorProto.COMPLEX64"
+    ),
+    int(TensorProto.COMPLEX128): TensorDtypeMap(
+        np.dtype("complex128"),
+        int(TensorProto.DOUBLE),
+        "TensorProto.COMPLEX128",
+    ),
+    int(TensorProto.UINT32): TensorDtypeMap(
+        np.dtype("uint32"), int(TensorProto.UINT64), "TensorProto.UINT32"
+    ),
+    int(TensorProto.UINT64): TensorDtypeMap(
+        np.dtype("uint64"), int(TensorProto.UINT64), "TensorProto.UINT64"
+    ),
+    int(TensorProto.STRING): TensorDtypeMap(
+        np.dtype("object"), int(TensorProto.STRING), "TensorProto.STRING"
+    ),
+    int(TensorProto.FLOAT8E4M3FN): TensorDtypeMap(
+        np.dtype(ml_dtypes.float8_e4m3fn),
+        int(TensorProto.INT32),
+        "TensorProto.FLOAT8E4M3FN",
+    ),
+    int(TensorProto.FLOAT8E4M3FNUZ): TensorDtypeMap(
+        np.dtype(ml_dtypes.float8_e4m3fnuz),
+        int(TensorProto.INT32),
+        "TensorProto.FLOAT8E4M3FNUZ",
+    ),
+    int(TensorProto.FLOAT8E5M2): TensorDtypeMap(
+        np.dtype(ml_dtypes.float8_e5m2),
+        int(TensorProto.INT32),
+        "TensorProto.FLOAT8E5M2",
+    ),
+    int(TensorProto.FLOAT8E5M2FNUZ): TensorDtypeMap(
+        np.dtype(ml_dtypes.float8_e5m2fnuz),
+        int(TensorProto.INT32),
+        "TensorProto.FLOAT8E5M2FNUZ",
+    ),
+    int(TensorProto.UINT4): TensorDtypeMap(
+        np.dtype(ml_dtypes.uint4), int(TensorProto.INT32), "TensorProto.UINT4"
+    ),
+    int(TensorProto.INT4): TensorDtypeMap(
+        np.dtype(ml_dtypes.int4), int(TensorProto.INT32), "TensorProto.INT4"
+    ),
+    int(TensorProto.FLOAT4E2M1): TensorDtypeMap(
+        np.dtype(ml_dtypes.float4_e2m1fn),
+        int(TensorProto.INT32),
+        "TensorProto.FLOAT4E2M1",
+    ),
+    int(TensorProto.FLOAT8E8M0): TensorDtypeMap(
+        np.dtype(ml_dtypes.float8_e8m0fnu),
+        int(TensorProto.INT32),
+        "TensorProto.FLOAT8E8M0",
+    ),
+}

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

@@ -18,6 +18,7 @@ from onnx.onnx_pb import (
     GraphProto,
     ModelProto,
     NodeProto,
+    OperatorSetIdProto,
     TensorProto,
     TypeProto,
 )
@@ -128,11 +129,25 @@ def function_expand_helper(
 
 
 def function_testcase_helper(
-    node: NodeProto, input_types: list[TypeProto], name: str
+    node: NodeProto,
+    input_types: list[TypeProto],
+    name: str,
+    opset_imports: Sequence[OperatorSetIdProto] | None = None,
 ) -> tuple[list[tuple[list[NodeProto], Any]], int]:
     test_op = node.op_type
     op_prefix = test_op + "_" + name + "_expanded_function_"
-    schema = onnx.defs.get_schema(test_op, domain=node.domain)
+    if opset_imports is None:
+        # No opset in the model. We take the most recent definition.
+        schema = onnx.defs.get_schema(test_op, domain=node.domain)
+    else:
+        # We take the function coming defined in the specific version mentioned
+        # in the model.
+        if len(opset_imports) != 1:
+            raise ValueError(
+                f"Only one domain is allowed but {len(opset_imports)} found."
+            )
+        version = opset_imports[0].version
+        schema = onnx.defs.get_schema(test_op, version, domain=node.domain)
 
     # an op schema may have several functions, each for one opset version
     # opset versions include the op's since_version and other opset versions
@@ -327,7 +342,9 @@ def expect(
     (
         expanded_tests,
         since_version,
-    ) = function_testcase_helper(node, merged_types, name)
+    ) = function_testcase_helper(
+        node, merged_types, name, opset_imports=kwargs.get("opset_imports")
+    )
     for expanded_function_nodes, func_opset_import in expanded_tests:
         kwargs["producer_name"] = "backend-test"
 

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

@@ -31,7 +31,7 @@ def argmax_use_numpy_select_last_index(
 class ArgMax(Base):
     @staticmethod
     def export_no_keepdims() -> None:
-        data = np.array([[2, 1], [3, 10]], dtype=np.float32)
+        data = np.array([[2, 2], [3, 10]], dtype=np.float32)
         axis = 1
         keepdims = 0
         node = onnx.helper.make_node(
@@ -55,7 +55,7 @@ class ArgMax(Base):
 
     @staticmethod
     def export_keepdims() -> None:
-        data = np.array([[2, 1], [3, 10]], dtype=np.float32)
+        data = np.array([[2, 2], [3, 10]], dtype=np.float32)
         axis = 1
         keepdims = 1
         node = onnx.helper.make_node(
@@ -76,7 +76,7 @@ class ArgMax(Base):
 
     @staticmethod
     def export_default_axes_keepdims() -> None:
-        data = np.array([[2, 1], [3, 10]], dtype=np.float32)
+        data = np.array([[2, 2], [3, 10]], dtype=np.float32)
         keepdims = 1
         node = onnx.helper.make_node(
             "ArgMax", inputs=["data"], outputs=["result"], keepdims=keepdims
@@ -103,7 +103,7 @@ class ArgMax(Base):
 
     @staticmethod
     def export_negative_axis_keepdims() -> None:
-        data = np.array([[2, 1], [3, 10]], dtype=np.float32)
+        data = np.array([[2, 2], [3, 10]], dtype=np.float32)
         axis = -1
         keepdims = 1
         node = onnx.helper.make_node(