onnx 1.16.2__cp310-cp310-win32.whl → 1.17.0__cp310-cp310-win32.whl

onnx/numpy_helper.py

@@ -1,24 +1,25 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 import sys
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Sequence
 
 import numpy as np
 
+import onnx._custom_element_types as custom_np_types
 from onnx import MapProto, OptionalProto, SequenceProto, TensorProto, helper, subbyte
 from onnx.external_data_helper import load_external_data_for_tensor, uses_external_data
 
 
-def combine_pairs_to_complex(fa: Sequence[int]) -> List[complex]:
+def combine_pairs_to_complex(fa: Sequence[int]) -> list[complex]:
     return [complex(fa[i * 2], fa[i * 2 + 1]) for i in range(len(fa) // 2)]
 
 
 def bfloat16_to_float32(
-    data: Union[np.int16, np.int32, np.ndarray],
-    dims: Optional[Union[int, Sequence[int]]] = None,
+    data: np.int16 | np.int32 | np.ndarray,
+    dims: int | Sequence[int] | None = None,
 ) -> np.ndarray:
     """Converts ndarray of bf16 (as uint32) to f32 (as uint32).
 
@@ -55,6 +56,7 @@ def _float8e4m3_to_float32_scalar(ival: int, fn: bool, uz: bool) -> np.float32:
         if ival == 127:  # noqa: PLR2004
             return np.float32(np.nan)
 
+    ival = np.uint32(ival)  # type: ignore[assignment]
     expo = (ival & 0x78) >> 3
     mant = ival & 0x07
     sign = ival & 0x80
@@ -86,8 +88,8 @@ _float8e4m3_to_float32 = np.vectorize(
 
 
 def float8e4m3_to_float32(
-    data: Union[np.int16, np.int32, np.ndarray],
-    dims: Optional[Union[int, Sequence[int]]] = None,
+    data: np.int16 | np.int32 | np.ndarray,
+    dims: int | Sequence[int] | None = None,
     fn: bool = True,
     uz: bool = False,
 ) -> np.ndarray:
@@ -133,6 +135,7 @@ def _float8e5m2_to_float32_scalar(ival: int, fn: bool, uz: bool) -> np.float32:
     else:
         raise NotImplementedError("fn and uz must be both False or True.")
 
+    ival = np.uint32(ival)  # type: ignore[assignment]
     expo = (ival & 0x7C) >> 2
     mant = ival & 0x03
     sign = ival & 0x80
@@ -160,8 +163,8 @@ _float8e5m2_to_float32 = np.vectorize(
 
 
 def float8e5m2_to_float32(
-    data: Union[np.int16, np.int32, np.ndarray],
-    dims: Optional[Union[int, Sequence[int]]] = None,
+    data: np.int16 | np.int32 | np.ndarray,
+    dims: int | Sequence[int] | None = None,
     fn: bool = False,
     uz: bool = False,
 ) -> np.ndarray:
@@ -186,8 +189,8 @@ def float8e5m2_to_float32(
 
 
 def unpack_int4(
-    data: Union[np.int32, np.ndarray],
-    dims: Union[int, Sequence[int]],
+    data: np.int32 | np.ndarray,
+    dims: int | Sequence[int],
     signed: bool,
 ) -> np.ndarray:
     """Converts ndarray of int4 (as packed uint8) to f32
@@ -218,7 +221,7 @@ def unpack_int4(
     return res
 
 
-def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PLR0911
+def _to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PLR0911
     """Converts a tensor def object to a numpy array.
 
     Args:
@@ -332,7 +335,79 @@ def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PL
     return np.asarray(data, dtype=storage_np_dtype).astype(np_dtype).reshape(dims)
 
 
-def from_array(arr: np.ndarray, name: Optional[str] = None) -> TensorProto:
+def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:
+    """Converts a tensor def object to a numpy array.
+    Supports types defined in :mod:`onnx._custom_element_types`.
+
+    Args:
+        tensor: a TensorProto object.
+        base_dir: if external tensor exists, base_dir can help to find the path to it
+
+    Returns:
+        arr: the converted array.
+    """
+    elem_type = tensor.data_type
+    if elem_type == TensorProto.BFLOAT16:
+        if uses_external_data(tensor):
+            load_external_data_for_tensor(tensor, base_dir)
+        data = tensor.int32_data
+        shape = tuple(tensor.dims)
+        y = np.empty(shape, dtype=custom_np_types.bfloat16).ravel()
+        for i, d in enumerate(data):
+            y[i] = d
+        return y.reshape(shape)
+
+    if elem_type in (
+        TensorProto.FLOAT8E4M3FN,
+        TensorProto.FLOAT8E4M3FNUZ,
+        TensorProto.FLOAT8E5M2,
+        TensorProto.FLOAT8E5M2FNUZ,
+    ):
+        m = {
+            TensorProto.FLOAT8E4M3FN: custom_np_types.float8e4m3fn,
+            TensorProto.FLOAT8E4M3FNUZ: custom_np_types.float8e4m3fnuz,
+            TensorProto.FLOAT8E5M2: custom_np_types.float8e5m2,
+            TensorProto.FLOAT8E5M2FNUZ: custom_np_types.float8e5m2fnuz,
+        }
+
+        if uses_external_data(tensor):
+            load_external_data_for_tensor(tensor, base_dir)
+        if tensor.HasField("raw_data"):
+            data = tensor.raw_data  # type: ignore[assignment]
+        else:
+            data = tensor.int32_data
+        shape = tuple(tensor.dims)
+        y = np.empty(shape, dtype=m[elem_type]).ravel()  # type: ignore[index]
+        for i, d in enumerate(data):
+            y[i] = d
+        return y.reshape(shape)
+
+    if elem_type in (TensorProto.UINT4, TensorProto.INT4):
+        if uses_external_data(tensor):
+            load_external_data_for_tensor(tensor, base_dir)
+        if tensor.HasField("raw_data"):
+            data = tensor.raw_data  # type: ignore[assignment]
+        else:
+            data = tensor.int32_data
+        shape = tuple(tensor.dims)
+        m = {
+            TensorProto.INT4: custom_np_types.int4,
+            TensorProto.UINT4: custom_np_types.uint4,
+        }
+        dtype = m[elem_type]  # type: ignore[index]
+        signed = elem_type == TensorProto.INT4
+        # 2 packed int4 elements must be represented as a single uint8 value.
+        # Therefore, y is np.uint8 (not the dtype to which the int4 maps)
+        y = np.empty(len(data), dtype=np.uint8).ravel()  # type: ignore[assignment]
+        for i, d in enumerate(data):
+            y[i] = d
+        unpacked_data = unpack_int4(y, dims=shape, signed=signed)
+        return unpacked_data.astype(dtype)
+
+    return _to_array(tensor, base_dir=base_dir)
+
+
+def _from_array(arr: np.ndarray, name: str | None = None) -> TensorProto:
     """Converts a numpy array to a tensor def.
 
     Args:
@@ -352,7 +427,7 @@ def from_array(arr: np.ndarray, name: Optional[str] = None) -> TensorProto:
     if name:
         tensor.name = name
 
-    if arr.dtype == object:
+    if arr.dtype == object or np.issubdtype(arr.dtype, np.str_):
         # Special care for strings.
         tensor.data_type = helper.np_dtype_to_tensor_dtype(arr.dtype)
         # TODO: Introduce full string support.
@@ -397,7 +472,68 @@ def from_array(arr: np.ndarray, name: Optional[str] = None) -> TensorProto:
     return tensor
 
 
-def to_list(sequence: SequenceProto) -> List[Any]:
+def from_array(tensor: np.ndarray, name: str | None = None) -> TensorProto:
+    """Converts an array into a TensorProto including
+    supported type defined in :mod:`onnx._custom_element_types`.
+
+    Args:
+        tensor: a numpy array.
+        name: (optional) the name of the tensor.
+
+    Returns:
+        TensorProto: the converted tensor def.
+    """
+    if not isinstance(tensor, np.ndarray):
+        return _from_array(tensor, name)
+
+    dt = tensor.dtype
+    if dt == custom_np_types.float8e4m3fn and dt.descr[0][0] == "e4m3fn":
+        to = TensorProto.FLOAT8E4M3FN
+        dt_to = np.uint8  # type: ignore[assignment]
+    elif dt == custom_np_types.float8e4m3fnuz and dt.descr[0][0] == "e4m3fnuz":
+        to = TensorProto.FLOAT8E4M3FNUZ
+        dt_to = np.uint8  # type: ignore[assignment]
+    elif dt == custom_np_types.float8e5m2 and dt.descr[0][0] == "e5m2":
+        to = TensorProto.FLOAT8E5M2
+        dt_to = np.uint8  # type: ignore[assignment]
+    elif dt == custom_np_types.float8e5m2fnuz and dt.descr[0][0] == "e5m2fnuz":
+        to = TensorProto.FLOAT8E5M2FNUZ
+        dt_to = np.uint8  # type: ignore[assignment]
+    elif dt == custom_np_types.bfloat16 and dt.descr[0][0] == "bfloat16":
+        to = TensorProto.BFLOAT16
+        dt_to = np.uint16  # type: ignore[assignment]
+    elif dt == custom_np_types.int4 and dt.descr[0][0] == "int4":
+        to = TensorProto.INT4
+        dt_to = np.int8  # type: ignore[assignment]
+    elif dt == custom_np_types.uint4 and dt.descr[0][0] == "uint4":
+        to = TensorProto.UINT4
+        dt_to = np.uint8  # type: ignore[assignment]
+    else:
+        return _from_array(tensor, name)
+
+    if to in (TensorProto.UINT4, TensorProto.INT4):
+        value = tensor.astype(dt_to).ravel()
+        if value.size % 2 == 1:
+            raise ValueError(
+                f"The conversion of a tensor of INT4 or UINT4 requires an even size "
+                f"(shape={tensor.shape}). Every byte stores two INT4 or two UINT4."
+            )
+        buffer = (value[::2] & 0x0F) + (value[1::2] << 4)
+
+        pb = TensorProto()
+        pb.dims.extend(tensor.shape)
+        if name:
+            pb.name = name
+        pb.raw_data = buffer.tobytes()
+        pb.data_type = to
+        return pb
+
+    t = _from_array(tensor.astype(dt_to), name)
+    t.data_type = to
+    return t
+
+
+def to_list(sequence: SequenceProto) -> list[Any]:
     """Converts a sequence def to a Python list.
 
     Args:
@@ -419,7 +555,7 @@ def to_list(sequence: SequenceProto) -> List[Any]:
 
 
 def from_list(
-    lst: List[Any], name: Optional[str] = None, dtype: Optional[int] = None
+    lst: list[Any], name: str | None = None, dtype: int | None = None
 ) -> SequenceProto:
     """Converts a list into a sequence def.
 
@@ -475,7 +611,7 @@ def from_list(
     return sequence
 
 
-def to_dict(map_proto: MapProto) -> Dict[Any, Any]:
+def to_dict(map_proto: MapProto) -> dict[Any, Any]:
     """Converts a map def to a Python dictionary.
 
     Args:
@@ -484,7 +620,7 @@ def to_dict(map_proto: MapProto) -> Dict[Any, Any]:
     Returns:
         The converted dictionary.
     """
-    key_list: List[Any] = []
+    key_list: list[Any] = []
     if map_proto.key_type == TensorProto.STRING:
         key_list = list(map_proto.string_keys)
     else:
@@ -501,7 +637,7 @@ def to_dict(map_proto: MapProto) -> Dict[Any, Any]:
     return dictionary
 
 
-def from_dict(dict_: Dict[Any, Any], name: Optional[str] = None) -> MapProto:
+def from_dict(dict_: dict[Any, Any], name: str | None = None) -> MapProto:
     """Converts a Python dictionary into a map def.
 
     Args:
@@ -559,7 +695,7 @@ def from_dict(dict_: Dict[Any, Any], name: Optional[str] = None) -> MapProto:
     return map_proto
 
 
-def to_optional(optional: OptionalProto) -> Optional[Any]:
+def to_optional(optional: OptionalProto) -> Any | None:
     """Converts an optional def to a Python optional.
 
     Args:
@@ -585,7 +721,7 @@ def to_optional(optional: OptionalProto) -> Optional[Any]:
 
 
 def from_optional(
-    opt: Optional[Any], name: Optional[str] = None, dtype: Optional[int] = None
+    opt: Any | None, name: str | None = None, dtype: int | None = None
 ) -> OptionalProto:
     """Converts an optional value into a Optional def.
 
@@ -637,7 +773,7 @@ def from_optional(
 
 
 def convert_endian(tensor: TensorProto) -> None:
-    """Call to convert endianess of raw data in tensor.
+    """Call to convert endianness of raw data in tensor.
 
     Args:
         tensor: TensorProto to be converted.
@@ -650,7 +786,7 @@ def convert_endian(tensor: TensorProto) -> None:
 
 
 def create_random_int(
-    input_shape: Tuple[int], dtype: np.dtype, seed: int = 1
+    input_shape: tuple[int], dtype: np.dtype, seed: int = 1
 ) -> np.ndarray:
     """Create random integer array for backend/test/case/node.
 

onnx/onnx-ml.proto

@@ -212,7 +212,7 @@ message NodeProto {
   repeated string output = 2;   // namespace Value
 
   // An optional identifier for this node in a graph.
-  // This field MAY be absent in ths version of the IR.
+  // This field MAY be absent in this version of the IR.
   optional string name = 3;     // namespace Node
 
   // The symbolic identifier of the Operator to execute.

onnx/onnx.in.proto

@@ -209,7 +209,7 @@ message NodeProto {
   repeated string output = 2;   // namespace Value
 
   // An optional identifier for this node in a graph.
-  // This field MAY be absent in ths version of the IR.
+  // This field MAY be absent in this version of the IR.
   optional string name = 3;     // namespace Node
 
   // The symbolic identifier of the Operator to execute.

onnx/onnx.proto

@@ -210,7 +210,7 @@ message NodeProto {
   repeated string output = 2;   // namespace Value
 
   // An optional identifier for this node in a graph.
-  // This field MAY be absent in ths version of the IR.
+  // This field MAY be absent in this version of the IR.
   optional string name = 3;     // namespace Node
 
   // The symbolic identifier of the Operator to execute.

onnx/onnx_cpp2py_export/defs.pyi

@@ -1,6 +1,4 @@
 """Submodule containing all the ONNX schema definitions."""
-from __future__ import annotations
-
 from typing import Sequence, overload
 
 from onnx import AttributeProto, FunctionProto

onnx/onnx_cpp2py_export/inliner.pyi

@@ -2,13 +2,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from __future__ import annotations
-
 def inline_local_functions(model: bytes, convert_version: bool) -> bytes:
     """Inlines calls to model-local function in input model and returns it.
     Both input and output are serialized ModelProtos.
     """
-    ...
 
 def inline_selected_functions(model: bytes, function_ids: list[tuple[str,str]], exclude: bool) -> bytes:
     """Inlines calls to selected model-local functions in input model and returns it.
@@ -16,4 +13,3 @@ def inline_selected_functions(model: bytes, function_ids: list[tuple[str,str]],
     case it inlines all functions except those specified in function_ids.
     Both input and output are serialized ModelProtos.
     """
-    ...

onnx/onnx_cpp2py_export/parser.pyi

@@ -5,7 +5,6 @@ def parse_model(model: str) -> tuple[bool, bytes, bytes]:
     If success-flag is true, then serialized-proto contains the parsed ModelProto.
     Otherwise, error-message contains a string describing the parse error.
     """
-    ...
 
 def parse_graph(graph: str) -> tuple[bool, bytes, bytes]:
     """Returns (success-flag, error-message, serialized-proto).
@@ -13,7 +12,6 @@ def parse_graph(graph: str) -> tuple[bool, bytes, bytes]:
     If success-flag is true, then serialized-proto contains the parsed GraphProto.
     Otherwise, error-message contains a string describing the parse error.
     """
-    ...
 
 def parse_function(function: str) -> tuple[bool, bytes, bytes]:
     """Returns (success-flag, error-message, serialized-proto).
@@ -21,7 +19,6 @@ def parse_function(function: str) -> tuple[bool, bytes, bytes]:
     If success-flag is true, then serialized-proto contains the parsed FunctionProto.
     Otherwise, error-message contains a string describing the parse error.
     """
-    ...
 
 def parse_node(node: str) -> tuple[bool, bytes, bytes]:
     """Returns (success-flag, error-message, serialized-proto).
@@ -29,4 +26,3 @@ def parse_node(node: str) -> tuple[bool, bytes, bytes]:
     If success-flag is true, then serialized-proto contains the parsed NodeProto.
     Otherwise, error-message contains a string describing the parse error.
     """
-    ...

onnx/parser.py

@@ -1,6 +1,7 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import onnx
 import onnx.onnx_cpp2py_export.parser as C  # noqa: N812

onnx/printer.py

@@ -1,14 +1,13 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
-
-from typing import Union
+from __future__ import annotations
 
 import onnx
 import onnx.onnx_cpp2py_export.printer as C  # noqa: N812
 
 
-def to_text(proto: Union[onnx.ModelProto, onnx.FunctionProto, onnx.GraphProto]) -> str:
+def to_text(proto: onnx.ModelProto | onnx.FunctionProto | onnx.GraphProto) -> str:
     if isinstance(proto, onnx.ModelProto):
         return C.model_to_text(proto.SerializeToString())
     if isinstance(proto, onnx.FunctionProto):

onnx/reference/__init__.py

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

onnx/reference/custom_element_types.py

@@ -1,13 +1,78 @@
 # Copyright (c) ONNX Project Contributors
-#
+
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import numpy as np
 
-bfloat16 = np.dtype((np.uint16, {"bfloat16": (np.uint16, 0)}))
-float8e4m3fn = np.dtype((np.uint8, {"e4m3fn": (np.uint8, 0)}))
-float8e4m3fnuz = np.dtype((np.uint8, {"e4m3fnuz": (np.uint8, 0)}))
-float8e5m2 = np.dtype((np.uint8, {"e5m2": (np.uint8, 0)}))
-float8e5m2fnuz = np.dtype((np.uint8, {"e5m2fnuz": (np.uint8, 0)}))
-uint4 = np.dtype((np.uint8, {"uint4": (np.uint8, 0)}))
-int4 = np.dtype((np.int8, {"int4": (np.int8, 0)}))
+try:
+    import ml_dtypes
+except ImportError:
+    ml_dtypes = None  # type: ignore[assignment]
+
+from onnx._custom_element_types import (
+    bfloat16,
+    float8e4m3fn,
+    float8e4m3fnuz,
+    float8e5m2,
+    float8e5m2fnuz,
+    int4,
+    uint4,
+)
+
+_supported_types = [
+    (bfloat16, "bfloat16", "bfloat16"),
+    (float8e4m3fn, "e4m3fn", "float8_e4m3fn"),
+    (float8e4m3fnuz, "e4m3fnuz", "float8_e4m3fnuz"),
+    (float8e5m2, "e5m2", "float8_e5m2"),
+    (float8e5m2fnuz, "e5m2fnuz", "float8_e5m2fnuz"),
+    (int4, "int4", "int4"),
+    (uint4, "uint4", "uint4"),
+]
+
+
+def convert_from_ml_dtypes(array: np.ndarray) -> np.ndarray:
+    """Detects the type and changes into one of the ONNX
+    defined custom types when ``ml_dtypes`` is installed.
+
+    Args:
+        array: Numpy array with a dtype from ml_dtypes.
+
+    Returns:
+        numpy array
+    """
+    if not ml_dtypes:
+        return array
+    for dtype, _, ml_name in _supported_types:
+        if array.dtype == getattr(ml_dtypes, ml_name):
+            return array.view(dtype=dtype)
+    return array
+
+
+def convert_to_ml_dtypes(array: np.ndarray) -> np.ndarray:
+    """Detects the type and changes into one of the type
+    defined in ``ml_dtypes`` if installed.
+
+    Args:
+        array: array
+
+    Returns:
+        numpy Numpy array with a dtype from ml_dtypes.
+    """
+    new_dt = None
+
+    for dtype, name, ml_name in _supported_types:
+        if array.dtype == dtype and array.dtype.descr[0][0] == name:
+            assert ml_dtypes, (
+                f"ml_dtypes is not installed and the tensor cannot "
+                f"be converted into ml_dtypes.{array.dtype.descr[0][0]}"
+            )
+            new_dt = getattr(ml_dtypes, ml_name)
+            break
+
+    if new_dt:
+        # int4, uint4, the representation uses 1 byte per element,
+        # only onnx storage uses 1 byte for two elements
+        return array.view(dtype=new_dt)
+
+    return array

onnx/reference/op_run.py

@@ -9,11 +9,7 @@ from typing import Any, ClassVar, Iterable
 import numpy as np
 
 from onnx import TensorProto
-from onnx.defs import get_all_schemas_with_history, get_schema, onnx_opset_version
-from onnx.helper import make_node, make_tensor_type_proto, np_dtype_to_tensor_dtype
-from onnx.numpy_helper import to_array, unpack_int4
-from onnx.onnx_pb import AttributeProto, GraphProto, NodeProto, TypeProto
-from onnx.reference.custom_element_types import (
+from onnx._custom_element_types import (
     bfloat16,
     float8e4m3fn,
     float8e4m3fnuz,
@@ -22,6 +18,10 @@ from onnx.reference.custom_element_types import (
     int4,
     uint4,
 )
+from onnx.defs import get_all_schemas_with_history, get_schema, onnx_opset_version
+from onnx.helper import make_node, make_tensor_type_proto, np_dtype_to_tensor_dtype
+from onnx.numpy_helper import to_array
+from onnx.onnx_pb import AttributeProto, GraphProto, NodeProto, TypeProto
 
 
 def _split_class_name(name):  # type: ignore
@@ -121,55 +121,7 @@ def to_sparse_tensor(att: AttributeProto) -> SparseTensor:
 
 
 def to_array_extended(tensor: TensorProto) -> np.ndarray:
-    """Similar to :func:`to_array` but deals with non-numpy types bfloat16,
-    float8e4m3fn, float8e4m3fnuz, float8e5m2, float8e5m2fnuz, uint4, int4.
-    """
-    elem_type = tensor.data_type
-    if elem_type == TensorProto.BFLOAT16:
-        data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        y = np.empty(shape, dtype=bfloat16).ravel()
-        for i, d in enumerate(data):
-            y[i] = d
-        return y.reshape(shape)
-
-    if elem_type in (
-        TensorProto.FLOAT8E4M3FN,
-        TensorProto.FLOAT8E4M3FNUZ,
-        TensorProto.FLOAT8E5M2,
-        TensorProto.FLOAT8E5M2FNUZ,
-    ):
-        m = {
-            TensorProto.FLOAT8E4M3FN: float8e4m3fn,
-            TensorProto.FLOAT8E4M3FNUZ: float8e4m3fnuz,
-            TensorProto.FLOAT8E5M2: float8e5m2,
-            TensorProto.FLOAT8E5M2FNUZ: float8e5m2fnuz,
-        }
-
-        if tensor.HasField("raw_data"):
-            data = tensor.raw_data  # type: ignore[assignment]
-        else:
-            data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        y = np.empty(shape, dtype=m[elem_type]).ravel()  # type: ignore[index]
-        for i, d in enumerate(data):
-            y[i] = d
-        return y.reshape(shape)
-    if elem_type in (TensorProto.UINT4, TensorProto.INT4):
-        if tensor.HasField("raw_data"):
-            data = tensor.raw_data  # type: ignore[assignment]
-        else:
-            data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        m = {TensorProto.INT4: int4, TensorProto.UINT4: uint4}
-        dtype = m[elem_type]  # type: ignore[index]
-        signed = elem_type == TensorProto.INT4
-        y = np.empty(len(data), dtype=dtype).ravel()
-        for i, d in enumerate(data):
-            y[i] = d
-
-        unpacked_data = unpack_int4(y, dims=shape, signed=signed)
-        return unpacked_data.astype(dtype)
+    """Alias for :func:`to_array`."""
     return to_array(tensor)
 
 
@@ -652,13 +604,16 @@ class OpRunExpand(OpRun):
     """Class any operator to avoid must inherit from."""
 
     def __init__(
-        self, onnx_node: NodeProto, run_params: dict[str, Any], impl: Any = None
+        self,
+        onnx_node: NodeProto,  # noqa: ARG002
+        run_params: dict[str, Any],  # noqa: ARG002
+        impl: Any = None,  # noqa: ARG002
     ):
         raise RuntimeError(
             f"The reference implementation must not use this node ({type(self)})."
         )
 
-    def _run(self, *inputs, **kwargs):
+    def _run(self, *inputs, **kwargs):  # noqa: ARG002
         raise RuntimeError(
             f"The reference implementation must not use this node ({type(self)})."
         )
@@ -735,7 +690,7 @@ class OpFunctionContextDependant(OpFunction):
         for t in inputs:
             try:
                 ttype = np_dtype_to_tensor_dtype(t.dtype)
-            except KeyError as e:
+            except KeyError:
                 if t.dtype == float8e4m3fn:
                     ttype = TensorProto.FLOAT8E4M3FN  # type: ignore[attr-defined]
                 elif t.dtype == float8e4m3fnuz:
@@ -751,7 +706,7 @@ class OpFunctionContextDependant(OpFunction):
                 elif t.dtype == int4:
                     ttype = TensorProto.INT4  # type: ignore[attr-defined]
                 else:
-                    raise e
+                    raise
             types.append(make_tensor_type_proto(ttype, t.shape))
         cl = self.parent._load_impl(self.onnx_node, types)
         inst = cl(self.onnx_node, self.run_params)

onnx/reference/ops/__init__.py

@@ -3,5 +3,6 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 from onnx.reference.ops._op_list import load_op

onnx/reference/ops/_helpers.py

@@ -3,16 +3,17 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
-from typing import Any, Dict, Union
+from typing import Any
 
 from onnx.reference.op_run import OpRun, _split_class_name
 
 
 def build_registered_operators_any_domain(
-    module_context: Dict[str, Any]
-) -> Dict[str, Dict[Union[int, None], OpRun]]:
-    reg_ops: Dict[str, Dict[Union[int, None], OpRun]] = {}
+    module_context: dict[str, Any],
+) -> dict[str, dict[int | None, OpRun]]:
+    reg_ops: dict[str, dict[int | None, OpRun]] = {}
     for class_name, class_type in module_context.items():
         if class_name.startswith("_") or class_name in {
             "Any",
@@ -20,6 +21,7 @@ def build_registered_operators_any_domain(
             "List",
             "TOptional",
             "Union",
+            "annotations",
             "cl",
             "class_name",
             "get_schema",