onnx 1.17.0__cp39-cp39-macosx_12_0_universal2.whl → 1.18.0__cp39-cp39-macosx_12_0_universal2.whl

onnx/helper.py

@@ -7,21 +7,11 @@ import collections.abc
 import numbers
 import struct
 from cmath import isnan
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    KeysView,
-    List,
-    Sequence,
-    Tuple,
-    TypeVar,
-    Union,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, Callable, TypeVar, Union, cast
 
 import google.protobuf.message
 import numpy as np
+import typing_extensions
 
 import onnx._custom_element_types as custom_np_types
 from onnx import (
@@ -46,9 +36,14 @@ from onnx import (
     subbyte,
 )
 
-VersionRowType = Union[Tuple[str, int, int, int], Tuple[str, int, int, int, int]]
-VersionTableType = List[VersionRowType]
-AssignmentBindingType = List[Tuple[str, str]]
+if TYPE_CHECKING:
+    from collections.abc import KeysView, Sequence
+
+    from google.protobuf.internal.containers import RepeatedCompositeFieldContainer
+
+VersionRowType = Union[tuple[str, int, int, int], tuple[str, int, int, int, int]]
+VersionTableType = list[VersionRowType]
+AssignmentBindingType = list[tuple[str, str]]
 
 # This is a copy of the documented version in https://github.com/onnx/onnx/blob/main/docs/Versioning.md#released-versions
 # Both must be updated whenever a new version of ONNX is released.
@@ -78,9 +73,10 @@ VERSION_TABLE: VersionTableType = [
     ("1.15.0", 9, 20, 4, 1),
     ("1.16.0", 10, 21, 5, 1),
     ("1.17.0", 10, 22, 5, 1),
+    ("1.18.0", 11, 23, 5, 1),
 ]
 
-VersionMapType = Dict[Tuple[str, int], int]
+VersionMapType = dict[tuple[str, int], int]
 
 
 def create_op_set_id_version_map(table: VersionTableType) -> VersionMapType:
@@ -307,8 +303,9 @@ def make_model(graph: GraphProto, **kwargs: Any) -> ModelProto:
     model.ir_version = IR_VERSION
     model.graph.CopyFrom(graph)
 
-    opset_imports: Sequence[OperatorSetIdProto] | None = None
-    opset_imports = kwargs.pop("opset_imports", None)  # type: ignore
+    opset_imports: Sequence[OperatorSetIdProto] | None = kwargs.pop(
+        "opset_imports", None
+    )
     if opset_imports is not None:
         model.opset_import.extend(opset_imports)
     else:
@@ -316,8 +313,7 @@ def make_model(graph: GraphProto, **kwargs: Any) -> ModelProto:
         imp = model.opset_import.add()
         imp.version = defs.onnx_opset_version()
 
-    functions: Sequence[FunctionProto] | None = None
-    functions = kwargs.pop("functions", None)  # type: ignore
+    functions: Sequence[FunctionProto] | None = kwargs.pop("functions", None)
     if functions is not None:
         model.functions.extend(functions)
 
@@ -360,20 +356,28 @@ 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]:
+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,
+)
+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:
     ival = int.from_bytes(struct.pack("<f", fval), "little")
     if truncate:
         return ival >> 16
@@ -386,7 +390,15 @@ def float32_to_bfloat16(fval: float, truncate: bool = False) -> int:
     return (ival + rounded) >> 16
 
 
-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,
+)
+def float32_to_float8e4m3(*args, **kwargs) -> int:
+    return _float32_to_float8e4m3(*args, **kwargs)
+
+
+def _float32_to_float8e4m3(  # noqa: PLR0911
     fval: float,
     scale: float = 1.0,
     fn: bool = True,
@@ -520,7 +532,15 @@ def float32_to_float8e4m3(  # noqa: PLR0911
         return int(ret)
 
 
-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,
+)
+def float32_to_float8e5m2(*args: Any, **kwargs: Any) -> int:
+    return _float32_to_float8e5m2(*args, **kwargs)
+
+
+def _float32_to_float8e5m2(  # noqa: PLR0911
     fval: float,
     scale: float = 1.0,
     fn: bool = False,
@@ -646,7 +666,15 @@ def float32_to_float8e5m2(  # noqa: PLR0911
         raise NotImplementedError("fn and uz must be both False or True.")
 
 
+@typing_extensions.deprecated(
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
+    category=DeprecationWarning,
+)
 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.
 
@@ -665,11 +693,43 @@ def pack_float32_to_4bit(array: np.ndarray | Sequence, signed: bool) -> np.ndarr
     if is_odd_volume:
         array_flat = np.append(array_flat, np.array([0]))
 
-    single_func = lambda x, y: subbyte.float32x2_to_4bitx2(x, y, signed)  # noqa: E731
+    def single_func(x, y) -> np.ndarray:
+        return subbyte._float32x2_to_4bitx2(x, y, signed)
+
     func = np.frompyfunc(single_func, 2, 1)
 
-    arr = func(array_flat[0::2], array_flat[1::2])
-    return arr.astype(np.uint8)  # type: ignore[no-any-return]
+    arr: np.ndarray = func(array_flat[0::2], array_flat[1::2])
+    return arr.astype(np.uint8)
+
+
+@typing_extensions.deprecated(
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
+    category=DeprecationWarning,
+)
+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.
+
+    Args:
+        array: array of float to convert and pack
+
+    Returns:
+        Packed array of float4e2m1 (as uint8) with size `ceil(farray.size/2)` (single dimension).
+    """
+    if not isinstance(array, np.ndarray):
+        array = np.asarray(array, dtype=np.float32)
+
+    array_flat = 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])
+    return arr.astype(np.uint8)
 
 
 def make_tensor(
@@ -702,7 +762,7 @@ def make_tensor(
     np_dtype = tensor_dtype_to_np_dtype(data_type)
 
     # Check number of vals specified equals tensor size
-    expected_size = 1
+    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:
@@ -714,13 +774,13 @@ def make_tensor(
             TensorProto.FLOAT8E5M2FNUZ,
         ):
             expected_size = 1
-        # NumPy doesn't have INT4. It is packed in couples to UINT8 buffers.
-        elif data_type in (TensorProto.UINT4, TensorProto.INT4):
-            expected_size = 0.5  # type: ignore[assignment]
+        # 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
         else:
             expected_size = np_dtype.itemsize
 
-    if type(vals) is np.ndarray and len(vals.shape) > 1:
+    if isinstance(vals, np.ndarray) and len(vals.shape) > 1:
         vals = vals.flatten()
     for d in dims:
         expected_size *= d
@@ -728,7 +788,7 @@ def make_tensor(
     if len(vals) != expected_size:
         # padding of half a byte is acceptable for 4bit types
         if not (
-            data_type in (TensorProto.UINT4, TensorProto.INT4)
+            data_type in (TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1)
             and len(vals) == expected_size + 0.5
         ):
             raise ValueError(
@@ -739,7 +799,7 @@ def make_tensor(
         tensor.raw_data = vals
     else:
         if data_type in (TensorProto.COMPLEX64, TensorProto.COMPLEX128):
-            vals = split_complex_to_pairs(vals)
+            vals = _split_complex_to_pairs(vals)
         elif data_type == TensorProto.FLOAT16:
             vals = (
                 np.array(vals).astype(np_dtype).view(dtype=np.uint16).flatten().tolist()
@@ -752,13 +812,13 @@ def make_tensor(
             TensorProto.FLOAT8E5M2FNUZ,
         ):
             fcast = {
-                TensorProto.BFLOAT16: float32_to_bfloat16,
-                TensorProto.FLOAT8E4M3FN: float32_to_float8e4m3,
-                TensorProto.FLOAT8E4M3FNUZ: lambda *args: float32_to_float8e4m3(  # type: ignore[misc]
+                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]
+                TensorProto.FLOAT8E5M2: _float32_to_float8e5m2,
+                TensorProto.FLOAT8E5M2FNUZ: lambda *args: _float32_to_float8e5m2(  # type: ignore[misc]
                     *args, fn=True, uz=True
                 ),
             }[
@@ -781,7 +841,9 @@ def make_tensor(
             # 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()
+            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:
@@ -824,16 +886,18 @@ def make_sequence(
 
     if elem_type == SequenceProto.UNDEFINED:
         return sequence
+
+    attribute: RepeatedCompositeFieldContainer | None = None
     if elem_type == SequenceProto.TENSOR:
         attribute = sequence.tensor_values
     elif elem_type == SequenceProto.SPARSE_TENSOR:
-        attribute = sequence.sparse_tensor_values  # type: ignore[assignment]
+        attribute = sequence.sparse_tensor_values
     elif elem_type == SequenceProto.SEQUENCE:
-        attribute = sequence.sequence_values  # type: ignore[assignment]
+        attribute = sequence.sequence_values
     elif elem_type == SequenceProto.MAP:
-        attribute = sequence.map_values  # type: ignore[assignment]
+        attribute = sequence.map_values
     elif elem_type == OptionalProto.OPTIONAL:
-        attribute = sequence.optional_values  # type: ignore[assignment]
+        attribute = sequence.optional_values
     else:
         raise TypeError("The element type in the input sequence is not supported.")
 
@@ -875,7 +939,7 @@ def make_map(
 def make_optional(
     name: str,
     elem_type: OptionalProto.DataType,
-    value: Any | None,
+    value: google.protobuf.message.Message | None,
 ) -> OptionalProto:
     """Make an Optional with specified value arguments."""
     optional = OptionalProto()
@@ -884,26 +948,30 @@ def make_optional(
 
     if elem_type == OptionalProto.UNDEFINED:
         return optional
+    attribute: google.protobuf.message.Message | None = None
     if elem_type == OptionalProto.TENSOR:
         attribute = optional.tensor_value
     elif elem_type == OptionalProto.SPARSE_TENSOR:
-        attribute = optional.sparse_tensor_value  # type: ignore[assignment]
+        attribute = optional.sparse_tensor_value
     elif elem_type == OptionalProto.SEQUENCE:
-        attribute = optional.sequence_value  # type: ignore[assignment]
+        attribute = optional.sequence_value
     elif elem_type == OptionalProto.MAP:
-        attribute = optional.map_value  # type: ignore[assignment]
+        attribute = optional.map_value
     elif elem_type == OptionalProto.OPTIONAL:
-        attribute = optional.optional_value  # type: ignore[assignment]
+        attribute = optional.optional_value
     else:
         raise TypeError("The element type in the input optional is not supported.")
 
+    assert value is not None
     attribute.CopyFrom(value)  # type: ignore[arg-type]
     return optional
 
 
 def _to_bytes(value: str | bytes) -> bytes:
     """Coerce a string (or bytes) value into UTF-8 bytes."""
-    return value if isinstance(value, bytes) else value.encode("utf-8")
+    if isinstance(value, str):
+        return value.encode("utf-8")
+    return value
 
 
 def make_attribute(
@@ -959,7 +1027,7 @@ def make_attribute(
                 (GraphProto, AttributeProto.GRAPHS),
                 (TypeProto, AttributeProto.TYPE_PROTOS),
             ):
-                if all(issubclass(t, exp_t) for t in types):  # type: ignore[arg-type]
+                if all(issubclass(t, exp_t) for t in types):
                     attr_type = exp_enum
                     break
             if attr_type is None:
@@ -1354,13 +1422,13 @@ def printable_dim(dim: TensorShapeProto.Dimension) -> str:
 
 def printable_type(t: TypeProto) -> str:
     if t.WhichOneof("value") == "tensor_type":
-        s = TensorProto.DataType.Name(t.tensor_type.elem_type)
+        s: str = TensorProto.DataType.Name(t.tensor_type.elem_type)
         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)))
             else:
                 s += ", scalar"
-        return s  # type: ignore[no-any-return]
+        return s
     if t.WhichOneof("value") is None:
         return ""
     return f"Unknown type {t.WhichOneof('value')}"
@@ -1438,9 +1506,7 @@ def printable_graph(graph: GraphProto, prefix: str = "") -> str:
     if len(graph.input):
         header.append("(")
         in_strs = []  # required inputs
-        in_with_init_strs = (
-            []
-        )  # optional inputs with initializer providing default value
+        in_with_init_strs: list = []  # optional inputs with initializer providing default value
         for inp in graph.input:
             if inp.name not in initializers:
                 in_strs.append(printable_value_info(inp))
@@ -1607,7 +1673,7 @@ def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> int:
         )
 
     if np.issubdtype(np_dtype, np.str_):
-        return TensorProto.STRING
+        return TensorProto.STRING  # type: ignore[no-any-return]
 
     if np_dtype in {
         custom_np_types.bfloat16,
@@ -1617,8 +1683,9 @@ def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> int:
         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]]
+        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."
@@ -1634,7 +1701,9 @@ def get_all_tensor_dtypes() -> KeysView[int]:
     return mapping.TENSOR_TYPE_MAP.keys()
 
 
-_ATTRIBUTE_TYPE_TO_STR = {k: v for v, k in AttributeProto.AttributeType.items()}
+_ATTRIBUTE_TYPE_TO_STR: dict[int, str] = {
+    k: v for v, k in AttributeProto.AttributeType.items()
+}
 
 
 def _attr_type_to_str(attr_type: int) -> str:
@@ -1647,5 +1716,5 @@ def _attr_type_to_str(attr_type: int) -> str:
         String representing the supplied attr_type.
     """
     if attr_type in AttributeProto.AttributeType.values():
-        return _ATTRIBUTE_TYPE_TO_STR[attr_type]  # type: ignore[no-any-return]
-    return AttributeProto.AttributeType.keys()[0]  # type: ignore[no-any-return]
+        return _ATTRIBUTE_TYPE_TO_STR[attr_type]
+    return AttributeProto.AttributeType.keys()[0]

onnx/hub.py

@@ -5,6 +5,7 @@
 
 This implements the python client for the ONNX model hub.
 """
+
 from __future__ import annotations
 
 import hashlib
@@ -13,7 +14,7 @@ import os
 import sys
 from io import BytesIO
 from os.path import join
-from typing import IO, Any, Dict, List, cast
+from typing import IO, Any, cast
 from urllib.error import HTTPError
 from urllib.request import urlopen
 
@@ -49,14 +50,14 @@ class ModelInfo:
         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.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.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.raw_model_info: dict[str, Any] = raw_model_info
@@ -363,10 +364,8 @@ def download_model_with_test_data(
                 "download the model from the model hub."
             )
 
-    # FIXME: Avoid index manipulation with magic numbers,
-    # remove ".tar.gz"
     local_model_with_data_dir_path = local_model_with_data_path[
-        0 : len(local_model_with_data_path) - 7
+        0 : len(local_model_with_data_path) - len(".tar.gz")
     ]
     onnx.utils._extract_model_safe(
         local_model_with_data_path, local_model_with_data_dir_path

onnx/inliner/inliner.cc

@@ -15,9 +15,9 @@
 
 #include "onnx/common/assertions.h"
 #include "onnx/common/constants.h"
-#include "onnx/common/interned_strings.h"
 #include "onnx/common/proto_util.h"
 #include "onnx/common/visitor.h"
+#include "onnx/defs/parser.h"
 #include "onnx/shape_inference/attribute_binder.h"
 #include "onnx/shape_inference/implementation.h"
 #include "onnx/version_converter/convert.h"
@@ -86,11 +86,11 @@ using RepeatedNodeProto = google::protobuf::RepeatedPtrField<NodeProto>;
 class NameGenerator : private Visitor {
  public:
   explicit NameGenerator(const GraphProto& graph) : index_(0) {
-    VisitGraph(graph);
+    NameGenerator::VisitGraph(graph);
   }
 
   explicit NameGenerator(const FunctionProto& function) : index_(0) {
-    VisitFunction(function);
+    NameGenerator::VisitFunction(function);
   }
 
   // Creates a new unique name, based on a suggested name, and adds it to the set
@@ -217,7 +217,7 @@ class InliningRenamer : private MutableVisitor {
     }
     for (; i < formals.size(); ++i) {
       std::string& formal = *formals.Mutable(i);
-      std::string rename_as = isOutput ? MakeUnique(formal) : std::string("");
+      std::string rename_as = isOutput ? MakeUnique(formal) : std::string();
       current_scope[formal] = rename_as;
       if (!rename_as.empty())
         formal = rename_as;
@@ -328,9 +328,9 @@ class ComputeInputs : private Visitor {
  public:
   std::vector<std::string> result;
 
-  ComputeInputs(const NodeProto& node) {
+  explicit ComputeInputs(const NodeProto& node) {
     result.reserve(node.input_size());
-    VisitNode(node);
+    ComputeInputs::VisitNode(node);
   }
 };
 
@@ -350,7 +350,7 @@ ConstNodeMap FindConstantNodes(const GraphProto& graph) {
   return result;
 }
 
-const TypeProto& GetType(const ModelProto& model, std::string var) {
+const TypeProto& GetType(const ModelProto& model, const std::string& var) {
   for (auto& vi : model.graph().value_info()) {
     if (vi.name() == var)
       return vi.type();
@@ -579,8 +579,7 @@ class VectorSet : public FunctionIdSet {
 std::unique_ptr<FunctionIdSet> ONNX_NAMESPACE::inliner::FunctionIdSet::Create(
     FunctionIdVector&& function_ids,
     bool invert) {
-  auto* p = new VectorSet(std::move(function_ids), invert);
-  return std::unique_ptr<FunctionIdSet>(p);
+  return std::make_unique<VectorSet>(std::move(function_ids), invert);
 }
 
 void InlineLocalFunctions(ModelProto& model, bool convert_version) {

onnx/mapping.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import Any, Dict, NamedTuple, Union, cast
+from typing import Any, NamedTuple, Union, cast
 
 import numpy as np
 
@@ -88,10 +88,13 @@ TENSOR_TYPE_MAP = {
     int(TensorProto.INT4): TensorDtypeMap(
         np.dtype("int8"), int(TensorProto.INT32), "TensorProto.INT4"
     ),
+    int(TensorProto.FLOAT4E2M1): TensorDtypeMap(
+        np.dtype("float32"), int(TensorProto.UINT8), "TensorProto.FLOAT4E2M1"
+    ),
 }
 
 
-class DeprecatedWarningDict(dict):  # type: ignore
+class DeprecatedWarningDict(dict):
     def __init__(
         self,
         dictionary: dict[int, int | str | np.dtype],
@@ -164,13 +167,14 @@ _NP_TYPE_TO_TENSOR_TYPE = {
         TensorProto.FLOAT8E5M2FNUZ,
         TensorProto.UINT4,
         TensorProto.INT4,
+        TensorProto.FLOAT4E2M1,
     )
 }
 
 # Currently native numpy does not support bfloat16 so TensorProto.BFLOAT16 is ignored for now
 # Numpy float32 array is only reversed to TensorProto.FLOAT
 NP_TYPE_TO_TENSOR_TYPE = DeprecatedWarningDict(
-    cast(Dict[int, Union[int, str, Any]], _NP_TYPE_TO_TENSOR_TYPE),
+    cast(dict[int, Union[int, str, Any]], _NP_TYPE_TO_TENSOR_TYPE),
     "NP_TYPE_TO_TENSOR_TYPE",
     "np_dtype_to_tensor_dtype",
 )
@@ -193,7 +197,7 @@ _STORAGE_TENSOR_TYPE_TO_FIELD = {
 }
 
 STORAGE_TENSOR_TYPE_TO_FIELD = DeprecatedWarningDict(
-    cast(Dict[int, Union[int, str, Any]], _STORAGE_TENSOR_TYPE_TO_FIELD),
+    cast(dict[int, Union[int, str, Any]], _STORAGE_TENSOR_TYPE_TO_FIELD),
     "STORAGE_TENSOR_TYPE_TO_FIELD",
 )
 

onnx/model_container.py

@@ -4,11 +4,12 @@
 """Implements function make_large_model to easily create and save models
 bigger than 2 Gb.
 """
+
 from __future__ import annotations
 
 import os
 import sys
-from typing import Any, Iterable
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -17,6 +18,9 @@ import onnx.external_data_helper as ext_data
 import onnx.helper
 import onnx.onnx_cpp2py_export.checker as c_checker
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
 
 def _set_external_data(
     tensor: onnx.TensorProto,
@@ -79,7 +83,7 @@ class ModelContainer:
     No tensor is stored on disk until the user explicitly saves the model.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.model_proto_: onnx.ModelProto | None = None
         self.large_initializers: dict[str, np.ndarray] = {}
 
@@ -126,13 +130,13 @@ class ModelContainer:
                 )
         self.large_initializers = large_initializers
 
-    def check_large_initializers(self):
+    def check_large_initializers(self) -> None:
         for tensor in ext_data._get_all_tensors(self.model_proto):
             if not ext_data.uses_external_data(tensor):
                 continue
             prop: onnx.StringStringEntryProto | None = None
-            for ext in tensor.external_data:  # type: ignore[assignment]
-                if ext.key == "location":  # type: ignore[attr-defined]
+            for ext in tensor.external_data:
+                if ext.key == "location":
                     prop = ext
             if prop is None:
                 raise RuntimeError(
@@ -195,9 +199,9 @@ class ModelContainer:
             if not ext_data.uses_external_data(tensor):
                 continue
             prop: onnx.StringStringEntryProto | None = None
-            for ext in tensor.external_data:  # type: ignore[assignment]
-                if ext.key == "location":  # type: ignore[attr-defined]
-                    prop = ext  # type: ignore[assignment]
+            for ext in tensor.external_data:
+                if ext.key == "location":
+                    prop = ext
             if prop is None:
                 raise RuntimeError(
                     f"No location found for tensor name {tensor.name!r}."