onnx 1.18.0__cp310-cp310-win_amd64.whl → 1.19.1rc1__cp310-cp310-win_amd64.whl

onnx/helper.py

@@ -4,36 +4,34 @@
 from __future__ import annotations
 
 import collections.abc
+import functools
+import math
 import numbers
 import struct
+import typing
 from cmath import isnan
-from typing import TYPE_CHECKING, Any, Callable, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, TypeVar, Union
 
 import google.protobuf.message
 import numpy as np
+import numpy.typing as npt
 import typing_extensions
 
-import onnx._custom_element_types as custom_np_types
-from onnx import (
-    IR_VERSION,
+import onnx
+from onnx import _mapping, defs, subbyte
+from onnx.onnx_data_pb import MapProto, OptionalProto, SequenceProto
+from onnx.onnx_pb import (
     AttributeProto,
     FunctionProto,
     GraphProto,
-    MapProto,
     ModelProto,
     NodeProto,
     OperatorSetIdProto,
-    OptionalProto,
-    SequenceProto,
-    SparseTensorProto,
     TensorProto,
     TensorShapeProto,
     TrainingInfoProto,
     TypeProto,
     ValueInfoProto,
-    defs,
-    mapping,
-    subbyte,
 )
 
 if TYPE_CHECKING:
@@ -74,12 +72,13 @@ VERSION_TABLE: VersionTableType = [
     ("1.16.0", 10, 21, 5, 1),
     ("1.17.0", 10, 22, 5, 1),
     ("1.18.0", 11, 23, 5, 1),
+    ("1.19.0", 12, 24, 5, 1),
 ]
 
 VersionMapType = dict[tuple[str, int], int]
 
 
-def create_op_set_id_version_map(table: VersionTableType) -> VersionMapType:
+def _create_op_set_id_version_map(table: VersionTableType) -> VersionMapType:
     """Create a map from (opset-domain, opset-version) to ir-version from above table."""
     result: VersionMapType = {}
 
@@ -96,7 +95,7 @@ def create_op_set_id_version_map(table: VersionTableType) -> VersionMapType:
     return result
 
 
-OP_SET_ID_VERSION_MAP = create_op_set_id_version_map(VERSION_TABLE)
+OP_SET_ID_VERSION_MAP = _create_op_set_id_version_map(VERSION_TABLE)
 
 
 def find_min_ir_version_for(
@@ -202,7 +201,7 @@ def make_graph(
     initializer: Sequence[TensorProto] | None = None,
     doc_string: str | None = None,
     value_info: Sequence[ValueInfoProto] | None = None,
-    sparse_initializer: Sequence[SparseTensorProto] | None = None,
+    sparse_initializer: Sequence[onnx.SparseTensorProto] | None = None,
 ) -> GraphProto:
     """Construct a GraphProto
 
@@ -214,7 +213,7 @@ def make_graph(
         initializer: list of TensorProto
         doc_string (string): graph documentation
         value_info: list of ValueInfoProto
-        sparse_initializer: list of SparseTensorProto
+        sparse_initializer: list of onnx.SparseTensorProto
     Returns:
         GraphProto
     """
@@ -300,7 +299,7 @@ def make_model(graph: GraphProto, **kwargs: Any) -> ModelProto:
     model = ModelProto()
     # Touch model.ir_version so it is stored as the version from which it is
     # generated.
-    model.ir_version = IR_VERSION
+    model.ir_version = onnx.IR_VERSION
     model.graph.CopyFrom(graph)
 
     opset_imports: Sequence[OperatorSetIdProto] | None = kwargs.pop(
@@ -356,28 +355,16 @@ def set_model_props(model: ModelProto, dict_value: dict[str, str]) -> None:
     set_metadata_props(model, dict_value)
 
 
-def _split_complex_to_pairs(ca: Sequence[np.complex64]) -> Sequence[int]:
-    return [
-        (ca[i // 2].real if (i % 2 == 0) else ca[i // 2].imag)  # type: ignore[misc]
-        for i in range(len(ca) * 2)
-    ]
-
-
 @typing_extensions.deprecated(
-    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
-    category=DeprecationWarning,
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion"
 )
-def float32_to_bfloat16(*args, **kwargs) -> int:
-    return _float32_to_bfloat16(*args, **kwargs)
-
-
-# convert a float32 value to a bfloat16 (as int)
-# By default, this conversion rounds-to-nearest-even and supports NaN
-# Setting `truncate` to True enables a simpler conversion. In this mode the
-# conversion is performed by simply dropping the 2 least significant bytes of
-# the significand. In this mode an error of up to 1 bit may be introduced and
-# preservation of NaN values is not be guaranteed.
-def _float32_to_bfloat16(fval: float, truncate: bool = False) -> int:
+def float32_to_bfloat16(fval: float, truncate: bool = False) -> int:
+    # convert a float32 value to a bfloat16 (as int)
+    # By default, this conversion rounds-to-nearest-even and supports NaN
+    # Setting `truncate` to True enables a simpler conversion. In this mode the
+    # conversion is performed by simply dropping the 2 least significant bytes of
+    # the significand. In this mode an error of up to 1 bit may be introduced and
+    # preservation of NaN values is not be guaranteed.
     ival = int.from_bytes(struct.pack("<f", fval), "little")
     if truncate:
         return ival >> 16
@@ -391,14 +378,9 @@ def _float32_to_bfloat16(fval: float, truncate: bool = False) -> int:
 
 
 @typing_extensions.deprecated(
-    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
-    category=DeprecationWarning,
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion"
 )
-def float32_to_float8e4m3(*args, **kwargs) -> int:
-    return _float32_to_float8e4m3(*args, **kwargs)
-
-
-def _float32_to_float8e4m3(  # noqa: PLR0911
+def float32_to_float8e4m3(  # noqa: PLR0911
     fval: float,
     scale: float = 1.0,
     fn: bool = True,
@@ -533,14 +515,9 @@ def _float32_to_float8e4m3(  # noqa: PLR0911
 
 
 @typing_extensions.deprecated(
-    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
-    category=DeprecationWarning,
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion"
 )
-def float32_to_float8e5m2(*args: Any, **kwargs: Any) -> int:
-    return _float32_to_float8e5m2(*args, **kwargs)
-
-
-def _float32_to_float8e5m2(  # noqa: PLR0911
+def float32_to_float8e5m2(  # noqa: PLR0911
     fval: float,
     scale: float = 1.0,
     fn: bool = False,
@@ -667,14 +644,9 @@ def _float32_to_float8e5m2(  # noqa: PLR0911
 
 
 @typing_extensions.deprecated(
-    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
-    category=DeprecationWarning,
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion"
 )
 def pack_float32_to_4bit(array: np.ndarray | Sequence, signed: bool) -> np.ndarray:
-    return _pack_float32_to_4bit(array, signed)
-
-
-def _pack_float32_to_4bit(array: np.ndarray | Sequence, signed: bool) -> np.ndarray:
     """Convert an array of float32 value to a 4bit data-type and pack every two concecutive elements in a byte.
     See :ref:`onnx-detail-int4` for technical details.
 
@@ -688,13 +660,13 @@ def _pack_float32_to_4bit(array: np.ndarray | Sequence, signed: bool) -> np.ndar
     if not isinstance(array, np.ndarray):
         array = np.asarray(array, dtype=np.float32)
 
-    array_flat = array.ravel()
+    array_flat: np.ndarray = array.ravel()
     is_odd_volume = np.prod(array.shape) % 2 == 1
     if is_odd_volume:
         array_flat = np.append(array_flat, np.array([0]))
 
     def single_func(x, y) -> np.ndarray:
-        return subbyte._float32x2_to_4bitx2(x, y, signed)
+        return subbyte.float32x2_to_4bitx2(x, y, signed)
 
     func = np.frompyfunc(single_func, 2, 1)
 
@@ -703,14 +675,9 @@ def _pack_float32_to_4bit(array: np.ndarray | Sequence, signed: bool) -> np.ndar
 
 
 @typing_extensions.deprecated(
-    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
-    category=DeprecationWarning,
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion"
 )
 def pack_float32_to_float4e2m1(array: np.ndarray | Sequence) -> np.ndarray:
-    return _pack_float32_to_float4e2m1(array)
-
-
-def _pack_float32_to_float4e2m1(array: np.ndarray | Sequence) -> np.ndarray:
     """Convert an array of float32 value to float4e2m1 and pack every two concecutive elements in a byte.
     See :ref:`onnx-detail-float4` for technical details.
 
@@ -723,17 +690,34 @@ def _pack_float32_to_float4e2m1(array: np.ndarray | Sequence) -> np.ndarray:
     if not isinstance(array, np.ndarray):
         array = np.asarray(array, dtype=np.float32)
 
-    array_flat = array.ravel()
+    array_flat: np.ndarray = array.ravel()
     is_odd_volume = np.prod(array.shape) % 2 == 1
     if is_odd_volume:
         array_flat = np.append(array_flat, np.array([0]))
 
-    arr = subbyte._float32x2_to_float4e2m1x2(array_flat[0::2], array_flat[1::2])
+    arr = subbyte.float32x2_to_float4e2m1x2(array_flat[0::2], array_flat[1::2])
     return arr.astype(np.uint8)
 
 
+def _pack_4bitx2(array: np.ndarray) -> npt.NDArray[np.uint8]:
+    """Convert a numpy array to flatten, packed int4/uint4. Elements must be in the correct range."""
+    # Create a 1D copy
+    array_flat = array.ravel().view(np.uint8).copy()
+    size = array.size
+    odd_sized = size % 2 == 1
+    if odd_sized:
+        array_flat.resize([size + 1], refcheck=False)
+    array_flat &= 0x0F
+    array_flat[1::2] <<= 4
+    return array_flat[0::2] | array_flat[1::2]  # type: ignore[return-type]
+
+
 def make_tensor(
-    name: str, data_type: int, dims: Sequence[int], vals: Any, raw: bool = False
+    name: str,
+    data_type: int,
+    dims: Sequence[int],
+    vals: Sequence[int | float] | bytes | np.ndarray,
+    raw: bool = False,
 ) -> TensorProto:
     """Make a TensorProto with specified arguments.  If raw is False, this
     function will choose the corresponding proto field to store the
@@ -742,12 +726,12 @@ def make_tensor(
     this case.
 
     Args:
-        name (string): tensor name
-        data_type (int): a value such as onnx.TensorProto.FLOAT
-        dims (List[int]): shape
+        name: tensor name
+        data_type: a value such as onnx.TensorProto.FLOAT
+        dims: shape
         vals: values
-        raw (bool): if True, vals contains the serialized content of the tensor,
-            otherwise, vals should be a list of values of the type defined by *data_type*
+        raw: if True, vals contains the serialized content of the tensor,
+            otherwise, vals should be a list of values of the type defined by ``data_type``.
 
     Returns:
         TensorProto
@@ -755,108 +739,85 @@ def make_tensor(
     tensor = TensorProto()
     tensor.data_type = data_type
     tensor.name = name
+    tensor.dims.extend(dims)
 
     if data_type == TensorProto.STRING and raw:
         raise TypeError("Can not use raw_data to store string type.")
 
     np_dtype = tensor_dtype_to_np_dtype(data_type)
 
-    # Check number of vals specified equals tensor size
-    expected_size: float = 1
     if raw:
-        # NumPy doesn't have BFLOAT16. TENSOR_TYPE_MAP maps it to float32, which has the wrong itemsize.
-        if data_type == TensorProto.BFLOAT16:
-            expected_size = 2
-        elif data_type in (
-            TensorProto.FLOAT8E4M3FN,
-            TensorProto.FLOAT8E4M3FNUZ,
-            TensorProto.FLOAT8E5M2,
-            TensorProto.FLOAT8E5M2FNUZ,
-        ):
-            expected_size = 1
         # NumPy doesn't have INT4/FP4. It is packed in couples to UINT8 buffers.
-        elif data_type in (TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1):
-            expected_size = 0.5
+        if data_type in {TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1}:
+            expected_size_bytes = 0.5
+        else:
+            expected_size_bytes = np_dtype.itemsize
+        expected_size_bytes *= math.prod(dims)
+        expected_size_bytes = math.ceil(expected_size_bytes)
+        if isinstance(vals, np.ndarray):
+            raw_data = vals.tobytes()
+        elif isinstance(vals, bytes):
+            raw_data = vals
         else:
-            expected_size = np_dtype.itemsize
-
-    if isinstance(vals, np.ndarray) and len(vals.shape) > 1:
-        vals = vals.flatten()
-    for d in dims:
-        expected_size *= d
-
-    if len(vals) != expected_size:
-        # padding of half a byte is acceptable for 4bit types
-        if not (
-            data_type in (TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1)
-            and len(vals) == expected_size + 0.5
-        ):
+            raise TypeError(
+                f"Raw data must be bytes or numpy.ndarray, but got {type(vals)}."
+            )
+        if len(raw_data) != expected_size_bytes:
             raise ValueError(
-                f"Number of values does not match tensor's size. Expected {expected_size}, but it is {len(vals)}. "
+                f"Raw data size does not match tensor's size. Expected {expected_size_bytes} bytes, but got {len(raw_data)} bytes."
             )
-
-    if raw:
-        tensor.raw_data = vals
+        tensor.raw_data = raw_data
+        return tensor
+
+    assert not raw, "Bug: raw should be False at this point."
+
+    if data_type == TensorProto.STRING:
+        vals = np.array(vals).flatten()
+        if len(vals) != 0:
+            vals = np.vectorize(_to_bytes)(vals)  # Convert to bytes
+    elif data_type in {
+        TensorProto.FLOAT8E4M3FN,
+        TensorProto.FLOAT8E4M3FNUZ,
+        TensorProto.FLOAT8E5M2,
+        TensorProto.FLOAT8E5M2FNUZ,
+    }:
+        # Float8 values are by default casted using saturating cast.
+        vals = onnx.numpy_helper.saturate_cast(np.asarray(vals), np_dtype).flatten()
+    elif data_type == TensorProto.FLOAT8E8M0:
+        vals = onnx.numpy_helper.to_float8e8m0(
+            np.asarray(vals), saturate=True, round_mode="up"
+        ).flatten()
     else:
-        if data_type in (TensorProto.COMPLEX64, TensorProto.COMPLEX128):
-            vals = _split_complex_to_pairs(vals)
-        elif data_type == TensorProto.FLOAT16:
-            vals = (
-                np.array(vals).astype(np_dtype).view(dtype=np.uint16).flatten().tolist()
-            )
-        elif data_type in (
-            TensorProto.BFLOAT16,
-            TensorProto.FLOAT8E4M3FN,
-            TensorProto.FLOAT8E4M3FNUZ,
-            TensorProto.FLOAT8E5M2,
-            TensorProto.FLOAT8E5M2FNUZ,
-        ):
-            fcast = {
-                TensorProto.BFLOAT16: _float32_to_bfloat16,
-                TensorProto.FLOAT8E4M3FN: _float32_to_float8e4m3,
-                TensorProto.FLOAT8E4M3FNUZ: lambda *args: _float32_to_float8e4m3(  # type: ignore[misc]
-                    *args, uz=True
-                ),
-                TensorProto.FLOAT8E5M2: _float32_to_float8e5m2,
-                TensorProto.FLOAT8E5M2FNUZ: lambda *args: _float32_to_float8e5m2(  # type: ignore[misc]
-                    *args, fn=True, uz=True
-                ),
-            }[
-                data_type  # type: ignore[index]
-            ]
-            vals = list(
-                map(  # type: ignore[call-overload]
-                    fcast,
-                    np.array(vals).astype(np_dtype).flatten().tolist(),
-                )
-            )
-        elif data_type in (
-            TensorProto.UINT4,
-            TensorProto.INT4,
-        ):
-            signed = data_type == TensorProto.INT4
-
-            # Two packed 4-bit values must be represented as a single uint8 value.
-            # Therefore, pack_float32_to_4bit() sets the dtype of the output vals
-            # to uint8 regardless of the value of 'signed'. Using int8 would cause
-            # the size of int4 tensors to increase ~5x if the tensor contains negative values (due to
-            # the way negative values are serialized by protobuf).
-            vals = _pack_float32_to_4bit(vals, signed=signed).flatten().tolist()
-        elif data_type == TensorProto.FLOAT4E2M1:
-            vals = _pack_float32_to_float4e2m1(vals).flatten().tolist()
-        elif data_type == TensorProto.BOOL:
-            vals = np.array(vals).astype(int)
-        elif data_type == TensorProto.STRING:
-            vals = np.array(vals).astype(bytes)
-        field = tensor_dtype_to_field(data_type)
-        getattr(tensor, field).extend(vals)
-    tensor.dims.extend(dims)
+        vals = np.asarray(vals, dtype=np_dtype).flatten()
+
+    if data_type == TensorProto.COMPLEX128:
+        vals = vals.view(np.float64)  # type: ignore[union-attr]
+    elif data_type == TensorProto.COMPLEX64:
+        vals = vals.view(np.float32)  # type: ignore[union-attr]
+    elif data_type in {TensorProto.BFLOAT16, TensorProto.FLOAT16}:
+        vals = vals.view(np.uint16)  # type: ignore[union-attr]
+    elif data_type in {
+        TensorProto.FLOAT8E4M3FN,
+        TensorProto.FLOAT8E4M3FNUZ,
+        TensorProto.FLOAT8E5M2,
+        TensorProto.FLOAT8E5M2FNUZ,
+        TensorProto.FLOAT8E8M0,
+    }:
+        vals = vals.view(np.uint8)  # type: ignore[union-attr]
+    elif data_type in {TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1}:
+        # Convert to packed 4-bit representation
+        vals = _pack_4bitx2(vals)  # type: ignore[union-attr,arg-type]
+    elif data_type == TensorProto.BOOL:
+        vals = vals.astype(np.uint8)  # type: ignore[union-attr]
+
+    field = tensor_dtype_to_field(data_type)
+    getattr(tensor, field).extend(vals)
     return tensor
 
 
 def make_sparse_tensor(
     values: TensorProto, indices: TensorProto, dims: Sequence[int]
-) -> SparseTensorProto:
+) -> onnx.SparseTensorProto:
     """Construct a SparseTensorProto
 
     Args:
@@ -867,7 +828,7 @@ def make_sparse_tensor(
     Returns:
         SparseTensorProto
     """
-    sparse = SparseTensorProto()
+    sparse = onnx.SparseTensorProto()
     sparse.values.CopyFrom(values)
     sparse.indices.CopyFrom(indices)
     sparse.dims.extend(dims)
@@ -1000,7 +961,7 @@ def make_attribute(
     elif isinstance(value, TensorProto):
         attr.t.CopyFrom(value)
         attr.type = AttributeProto.TENSOR
-    elif isinstance(value, SparseTensorProto):
+    elif isinstance(value, onnx.SparseTensorProto):
         attr.sparse_tensor.CopyFrom(value)
         attr.type = AttributeProto.SPARSE_TENSOR
     elif isinstance(value, GraphProto):
@@ -1023,11 +984,11 @@ def make_attribute(
                 (numbers.Real, AttributeProto.FLOATS),
                 ((str, bytes), AttributeProto.STRINGS),
                 (TensorProto, AttributeProto.TENSORS),
-                (SparseTensorProto, AttributeProto.SPARSE_TENSORS),
+                (onnx.SparseTensorProto, AttributeProto.SPARSE_TENSORS),
                 (GraphProto, AttributeProto.GRAPHS),
                 (TypeProto, AttributeProto.TYPE_PROTOS),
             ):
-                if all(issubclass(t, exp_t) for t in types):
+                if all(issubclass(t, exp_t) for t in types):  # type: ignore[arg-type]
                     attr_type = exp_enum
                     break
             if attr_type is None:
@@ -1422,7 +1383,7 @@ def printable_dim(dim: TensorShapeProto.Dimension) -> str:
 
 def printable_type(t: TypeProto) -> str:
     if t.WhichOneof("value") == "tensor_type":
-        s: str = TensorProto.DataType.Name(t.tensor_type.elem_type)
+        s: str = TensorProto.DataType.Name(t.tensor_type.elem_type)  # type: ignore[attr-defined]
         if t.tensor_type.HasField("shape"):
             if len(t.tensor_type.shape.dim):
                 s += str(", " + "x".join(map(printable_dim, t.tensor_type.shape.dim)))
@@ -1443,7 +1404,7 @@ def printable_value_info(v: ValueInfoProto) -> str:
 
 def printable_tensor_proto(t: TensorProto) -> str:
     s = f"%{t.name}["
-    s += TensorProto.DataType.Name(t.data_type)
+    s += TensorProto.DataType.Name(t.data_type)  # type: ignore[attr-defined]
     if t.dims is not None:
         if len(t.dims):
             s += str(", " + "x".join(map(str, t.dims)))
@@ -1488,9 +1449,15 @@ def printable_node(
     return prefix + " ".join(content)
 
 
+@typing_extensions.deprecated(
+    "Deprecated since 1.19. Consider using onnx.printer.to_text() instead."
+)
 def printable_graph(graph: GraphProto, prefix: str = "") -> str:
     """Display a GraphProto as a string.
 
+    .. deprecated:: 1.19
+        Consider using :func:`onnx.printer.to_text` instead.
+
     Args:
         graph (GraphProto): the graph to display
         prefix (string): prefix of every line
@@ -1616,7 +1583,7 @@ def tensor_dtype_to_np_dtype(tensor_dtype: int) -> np.dtype:
     Returns:
         numpy's data_type
     """
-    return mapping.TENSOR_TYPE_MAP[tensor_dtype].np_dtype
+    return _mapping.TENSOR_TYPE_MAP[tensor_dtype].np_dtype
 
 
 def tensor_dtype_to_storage_tensor_dtype(tensor_dtype: int) -> int:
@@ -1628,7 +1595,7 @@ def tensor_dtype_to_storage_tensor_dtype(tensor_dtype: int) -> int:
     Returns:
         data_type for storage
     """
-    return mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype
+    return _mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype
 
 
 def tensor_dtype_to_string(tensor_dtype: int) -> str:
@@ -1640,9 +1607,10 @@ def tensor_dtype_to_string(tensor_dtype: int) -> str:
     Returns:
         the name of data_type
     """
-    return mapping.TENSOR_TYPE_MAP[tensor_dtype].name
+    return _mapping.TENSOR_TYPE_MAP[tensor_dtype].name
 
 
+@functools.lru_cache(None)
 def tensor_dtype_to_field(tensor_dtype: int) -> str:
     """Convert a TensorProto's data_type to corresponding field name for storage. It can be used while making tensors.
 
@@ -1652,12 +1620,22 @@ def tensor_dtype_to_field(tensor_dtype: int) -> str:
     Returns:
         field name
     """
-    return mapping._STORAGE_TENSOR_TYPE_TO_FIELD[
-        mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype
+    storage_tensor_type_to_field = {
+        int(TensorProto.FLOAT): "float_data",
+        int(TensorProto.INT32): "int32_data",
+        int(TensorProto.INT64): "int64_data",
+        int(TensorProto.DOUBLE): "double_data",
+        int(TensorProto.UINT32): "uint64_data",
+        int(TensorProto.UINT64): "uint64_data",
+        int(TensorProto.STRING): "string_data",
+    }
+    return storage_tensor_type_to_field[
+        _mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype
     ]
 
 
-def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> int:
+@functools.lru_cache(None)
+def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> TensorProto.DataType:
     """Convert a numpy's dtype to corresponding tensor type. It can be used while converting numpy arrays to tensors.
 
     Args:
@@ -1666,27 +1644,14 @@ def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> int:
     Returns:
         TensorsProto's data_type
     """
-    if np_dtype in mapping._NP_TYPE_TO_TENSOR_TYPE:
-        return cast(
-            int,
-            mapping._NP_TYPE_TO_TENSOR_TYPE[np_dtype],
-        )
-
+    _np_dtype_to_tensor_dtype = {
+        v.np_dtype: k for k, v in _mapping.TENSOR_TYPE_MAP.items()
+    }
+    if np_dtype in _np_dtype_to_tensor_dtype:
+        return typing.cast("TensorProto.DataType", _np_dtype_to_tensor_dtype[np_dtype])
     if np.issubdtype(np_dtype, np.str_):
         return TensorProto.STRING  # type: ignore[no-any-return]
 
-    if np_dtype in {
-        custom_np_types.bfloat16,
-        custom_np_types.float8e4m3fn,
-        custom_np_types.float8e4m3fnuz,
-        custom_np_types.float8e5m2,
-        custom_np_types.float8e5m2fnuz,
-        custom_np_types.int4,
-        custom_np_types.uint4,
-        custom_np_types.float4e2m1,
-    }:
-        return custom_np_types.mapping_name_to_data_type[np_dtype.descr[0][0]]  # type: ignore[no-any-return]
-
     raise ValueError(
         f"Unable to convert type {np_dtype!r} into TensorProto element type."
     )
@@ -1698,11 +1663,12 @@ def get_all_tensor_dtypes() -> KeysView[int]:
     Returns:
         all tensor types from TensorProto
     """
-    return mapping.TENSOR_TYPE_MAP.keys()
+    return _mapping.TENSOR_TYPE_MAP.keys()
 
 
 _ATTRIBUTE_TYPE_TO_STR: dict[int, str] = {
-    k: v for v, k in AttributeProto.AttributeType.items()
+    k: v
+    for v, k in AttributeProto.AttributeType.items()  # type: ignore[attr-defined]
 }
 
 
@@ -1715,6 +1681,6 @@ def _attr_type_to_str(attr_type: int) -> str:
     Returns:
         String representing the supplied attr_type.
     """
-    if attr_type in AttributeProto.AttributeType.values():
+    if attr_type in AttributeProto.AttributeType.values():  # type: ignore[attr-defined]
         return _ATTRIBUTE_TYPE_TO_STR[attr_type]
-    return AttributeProto.AttributeType.keys()[0]
+    return AttributeProto.AttributeType.keys()[0]  # type: ignore

onnx/hub.py

@@ -47,19 +47,21 @@ class ModelInfo:
         Args:
             raw_model_info: A JSON dict containing the model info.
         """
-        self.model = cast(str, raw_model_info["model"])
+        self.model = cast("str", raw_model_info["model"])
 
-        self.model_path = cast(str, raw_model_info["model_path"])
-        self.metadata: dict[str, Any] = cast(dict[str, Any], raw_model_info["metadata"])
+        self.model_path = cast("str", raw_model_info["model_path"])
+        self.metadata: dict[str, Any] = cast(
+            "dict[str, Any]", raw_model_info["metadata"]
+        )
         self.model_sha: str | None = None
         if "model_sha" in self.metadata:
-            self.model_sha = cast(str, self.metadata["model_sha"])
+            self.model_sha = cast("str", self.metadata["model_sha"])
 
         self.tags: set[str] = set()
         if "tags" in self.metadata:
-            self.tags = set(cast(list[str], self.metadata["tags"]))
+            self.tags = set(cast("list[str]", self.metadata["tags"]))
 
-        self.opset = cast(int, raw_model_info["opset_version"])
+        self.opset = cast("int", raw_model_info["opset_version"])
         self.raw_model_info: dict[str, Any] = raw_model_info
 
     def __str__(self) -> str:
@@ -167,7 +169,7 @@ def list_models(
     try:
         with urlopen(manifest_url) as response:
             manifest: list[ModelInfo] = [
-                ModelInfo(info) for info in json.load(cast(IO[str], response))
+                ModelInfo(info) for info in json.load(cast("IO[str]", response))
             ]
     except HTTPError as e:
         raise AssertionError(f"Could not find manifest at {manifest_url}") from e
@@ -287,7 +289,7 @@ def load(
                 "download the model from the model hub."
             )
 
-    return onnx.load(cast(IO[bytes], BytesIO(model_bytes)))
+    return onnx.load(cast("IO[bytes]", BytesIO(model_bytes)))
 
 
 def download_model_with_test_data(