onnx 1.16.2__cp311-cp311-win32.whl → 1.18.0__cp311-cp311-win32.whl

onnx/helper.py

@@ -1,28 +1,19 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import collections.abc
 import numbers
 import struct
 from cmath import isnan
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    KeysView,
-    List,
-    Optional,
-    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 (
     IR_VERSION,
     AttributeProto,
@@ -45,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.
@@ -76,9 +72,11 @@ VERSION_TABLE: VersionTableType = [
     ("1.14.1", 9, 19, 3, 1),
     ("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:
@@ -116,7 +114,7 @@ def find_min_ir_version_for(
     """
     default_min_version = 3
 
-    def find_min(domain: Union[str, None], version: int) -> int:
+    def find_min(domain: str | None, version: int) -> int:
         key = (domain or "ai.onnx", version)
         if key in OP_SET_ID_VERSION_MAP:
             return OP_SET_ID_VERSION_MAP[key]
@@ -133,10 +131,10 @@ def make_node(
     op_type: str,
     inputs: Sequence[str],
     outputs: Sequence[str],
-    name: Optional[str] = None,
-    doc_string: Optional[str] = None,
-    domain: Optional[str] = None,
-    overload: Optional[str] = None,
+    name: str | None = None,
+    doc_string: str | None = None,
+    domain: str | None = None,
+    overload: str | None = None,
     **kwargs: Any,
 ) -> NodeProto:
     """Construct a NodeProto.
@@ -201,10 +199,10 @@ def make_graph(
     name: str,
     inputs: Sequence[ValueInfoProto],
     outputs: Sequence[ValueInfoProto],
-    initializer: Optional[Sequence[TensorProto]] = None,
-    doc_string: Optional[str] = None,
-    value_info: Optional[Sequence[ValueInfoProto]] = None,
-    sparse_initializer: Optional[Sequence[SparseTensorProto]] = None,
+    initializer: Sequence[TensorProto] | None = None,
+    doc_string: str | None = None,
+    value_info: Sequence[ValueInfoProto] | None = None,
+    sparse_initializer: Sequence[SparseTensorProto] | None = None,
 ) -> GraphProto:
     """Construct a GraphProto
 
@@ -261,11 +259,11 @@ def make_function(
     outputs: Sequence[str],
     nodes: Sequence[NodeProto],
     opset_imports: Sequence[OperatorSetIdProto],
-    attributes: Optional[Sequence[str]] = None,
-    attribute_protos: Optional[Sequence[AttributeProto]] = None,
-    doc_string: Optional[str] = None,
-    overload: Optional[str] = None,
-    value_info: Optional[Sequence[ValueInfoProto]] = None,
+    attributes: Sequence[str] | None = None,
+    attribute_protos: Sequence[AttributeProto] | None = None,
+    doc_string: str | None = None,
+    overload: str | None = None,
+    value_info: Sequence[ValueInfoProto] | None = None,
 ) -> FunctionProto:
     if attributes is None:
         attributes = []
@@ -305,8 +303,9 @@ def make_model(graph: GraphProto, **kwargs: Any) -> ModelProto:
     model.ir_version = IR_VERSION
     model.graph.CopyFrom(graph)
 
-    opset_imports: Optional[Sequence[OperatorSetIdProto]] = 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:
@@ -314,8 +313,7 @@ def make_model(graph: GraphProto, **kwargs: Any) -> ModelProto:
         imp = model.opset_import.add()
         imp.version = defs.onnx_opset_version()
 
-    functions: Optional[Sequence[FunctionProto]] = 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)
 
@@ -337,10 +335,15 @@ def make_model_gen_version(graph: GraphProto, **kwargs: Any) -> ModelProto:
 
 
 def set_metadata_props(
-    proto: Union[
-        ModelProto, GraphProto, FunctionProto, NodeProto, TensorProto, ValueInfoProto
-    ],
-    dict_value: Dict[str, str],
+    proto: (
+        ModelProto
+        | GraphProto
+        | FunctionProto
+        | NodeProto
+        | TensorProto
+        | ValueInfoProto
+    ),
+    dict_value: dict[str, str],
 ) -> None:
     del proto.metadata_props[:]
     for k, v in dict_value.items():
@@ -349,24 +352,32 @@ def set_metadata_props(
         entry.value = v
 
 
-def set_model_props(model: ModelProto, dict_value: Dict[str, str]) -> None:
+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
@@ -379,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,
@@ -513,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,
@@ -639,9 +666,15 @@ def float32_to_float8e5m2(  # noqa: PLR0911
         raise NotImplementedError("fn and uz must be both False or True.")
 
 
-def pack_float32_to_4bit(
-    array: Union[np.ndarray, Sequence], signed: bool
-) -> np.ndarray:
+@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.
 
@@ -660,11 +693,43 @@ def pack_float32_to_4bit(
     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(
@@ -697,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:
@@ -709,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
@@ -723,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(
@@ -734,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()
@@ -747,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
                 ),
             }[
@@ -770,12 +835,15 @@ def make_tensor(
             TensorProto.INT4,
         ):
             signed = data_type == TensorProto.INT4
-            vals = (
-                pack_float32_to_4bit(vals, signed=signed)
-                .astype(np_dtype)
-                .flatten()
-                .tolist()
-            )
+
+            # 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:
@@ -818,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.")
 
@@ -836,7 +906,7 @@ def make_sequence(
 
 
 def make_map(
-    name: str, key_type: int, keys: List[Any], values: SequenceProto
+    name: str, key_type: int, keys: list[Any], values: SequenceProto
 ) -> MapProto:
     """Make a Map with specified key-value pair arguments.
 
@@ -869,7 +939,7 @@ def make_map(
 def make_optional(
     name: str,
     elem_type: OptionalProto.DataType,
-    value: Optional[Any],
+    value: google.protobuf.message.Message | None,
 ) -> OptionalProto:
     """Make an Optional with specified value arguments."""
     optional = OptionalProto()
@@ -878,33 +948,37 @@ 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: Union[str, bytes]) -> bytes:
+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(
     key: str,
     value: Any,
-    doc_string: Optional[str] = None,
-    attr_type: Optional[int] = None,
+    doc_string: str | None = None,
+    attr_type: int | None = None,
 ) -> AttributeProto:
     """Makes an AttributeProto based on the value type."""
     attr = AttributeProto()
@@ -953,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:
@@ -995,7 +1069,7 @@ def make_attribute(
 
 
 def make_attribute_ref(
-    name: str, attr_type: AttributeProto.AttributeType, doc_string: Optional[str] = None
+    name: str, attr_type: AttributeProto.AttributeType, doc_string: str | None = None
 ) -> AttributeProto:
     """Make an AttributeProto holding a reference to the parent function's attribute of given name and type."""
     attr = AttributeProto()
@@ -1059,8 +1133,8 @@ def make_empty_tensor_value_info(name: str) -> ValueInfoProto:
 
 def make_tensor_type_proto(
     elem_type: int,
-    shape: Optional[Sequence[Union[str, int, None]]],
-    shape_denotation: Optional[List[str]] = None,
+    shape: Sequence[str | int | None] | None,
+    shape_denotation: list[str] | None = None,
 ) -> TypeProto:
     """Makes a Tensor TypeProto based on the data type and shape."""
     type_proto = TypeProto()
@@ -1105,9 +1179,9 @@ def make_tensor_type_proto(
 def make_tensor_value_info(
     name: str,
     elem_type: int,
-    shape: Optional[Sequence[Union[str, int, None]]],
+    shape: Sequence[str | int | None] | None,
     doc_string: str = "",
-    shape_denotation: Optional[List[str]] = None,
+    shape_denotation: list[str] | None = None,
 ) -> ValueInfoProto:
     """Makes a ValueInfoProto based on the data type and shape."""
     value_info_proto = ValueInfoProto()
@@ -1122,8 +1196,8 @@ def make_tensor_value_info(
 
 def make_sparse_tensor_type_proto(
     elem_type: int,
-    shape: Optional[Sequence[Union[str, int, None]]],
-    shape_denotation: Optional[List[str]] = None,
+    shape: Sequence[str | int | None] | None,
+    shape_denotation: list[str] | None = None,
 ) -> TypeProto:
     """Makes a SparseTensor TypeProto based on the data type and shape."""
     type_proto = TypeProto()
@@ -1168,9 +1242,9 @@ def make_sparse_tensor_type_proto(
 def make_sparse_tensor_value_info(
     name: str,
     elem_type: int,
-    shape: Optional[Sequence[Union[str, int, None]]],
+    shape: Sequence[str | int | None] | None,
     doc_string: str = "",
-    shape_denotation: Optional[List[str]] = None,
+    shape_denotation: list[str] | None = None,
 ) -> ValueInfoProto:
     """Makes a SparseTensor ValueInfoProto based on the data type and shape."""
     value_info_proto = ValueInfoProto()
@@ -1231,7 +1305,7 @@ def make_value_info(
     return value_info_proto
 
 
-def _sanitize_str(s: Union[str, bytes]) -> str:
+def _sanitize_str(s: str | bytes) -> str:
     if isinstance(s, str):
         sanitized = s
     elif isinstance(s, bytes):
@@ -1246,9 +1320,9 @@ def _sanitize_str(s: Union[str, bytes]) -> str:
 def make_tensor_sequence_value_info(
     name: str,
     elem_type: int,
-    shape: Optional[Sequence[Union[str, int, None]]],
+    shape: Sequence[str | int | None] | None,
     doc_string: str = "",
-    elem_shape_denotation: Optional[List[str]] = None,
+    elem_shape_denotation: list[str] | None = None,
 ) -> ValueInfoProto:
     """Makes a Sequence[Tensors] ValueInfoProto based on the data type and shape."""
     value_info_proto = ValueInfoProto()
@@ -1265,7 +1339,7 @@ def make_tensor_sequence_value_info(
 
 def printable_attribute(
     attr: AttributeProto, subgraphs: bool = False
-) -> Union[str, Tuple[str, List[GraphProto]]]:
+) -> str | tuple[str, list[GraphProto]]:
     content = []
     content.append(attr.name)
     content.append("=")
@@ -1348,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')}"
@@ -1381,13 +1455,13 @@ def printable_tensor_proto(t: TensorProto) -> str:
 
 def printable_node(
     node: NodeProto, prefix: str = "", subgraphs: bool = False
-) -> Union[str, Tuple[str, List[GraphProto]]]:
+) -> str | tuple[str, list[GraphProto]]:
     content = []
     if len(node.output):
         content.append(", ".join([f"%{name}" for name in node.output]))
         content.append("=")
     # To deal with nested graphs
-    graphs: List[GraphProto] = []
+    graphs: list[GraphProto] = []
     printed_attrs = []
     for attr in node.attribute:
         if subgraphs:
@@ -1432,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))
@@ -1444,7 +1516,7 @@ def printable_graph(graph: GraphProto, prefix: str = "") -> str:
             content.append(prefix + " ".join(header))
             header = []
             for line in in_strs:
-                content.append(prefix + "  " + line)
+                content.append(prefix + "  " + line)  # noqa: PERF401
         header.append(")")
 
         if in_with_init_strs:
@@ -1452,7 +1524,7 @@ def printable_graph(graph: GraphProto, prefix: str = "") -> str:
             content.append(prefix + " ".join(header))
             header = []
             for line in in_with_init_strs:
-                content.append(prefix + "  " + line)
+                content.append(prefix + "  " + line)  # noqa: PERF401
             header.append(")")
 
         # from IR 4 onwards an initializer is not required to have a matching graph input
@@ -1468,12 +1540,12 @@ def printable_graph(graph: GraphProto, prefix: str = "") -> str:
             content.append(prefix + " ".join(header))
             header = []
             for line in init_strs:
-                content.append(prefix + "  " + line)
+                content.append(prefix + "  " + line)  # noqa: PERF401
             header.append(")")
 
     header.append("{")
     content.append(prefix + " ".join(header))
-    graphs: List[GraphProto] = []
+    graphs: list[GraphProto] = []
     # body
     for node in graph.node:
         contents_subgraphs = printable_node(node, indent, subgraphs=True)
@@ -1489,7 +1561,7 @@ def printable_graph(graph: GraphProto, prefix: str = "") -> str:
     # closing bracket
     content.append(prefix + "}")
     for g in graphs:
-        content.append("\n" + printable_graph(g))
+        content.append("\n" + printable_graph(g))  # noqa: PERF401
     return "\n".join(content)
 
 
@@ -1513,8 +1585,8 @@ def strip_doc_string(proto: google.protobuf.message.Message) -> None:
 def make_training_info(
     algorithm: GraphProto,
     algorithm_bindings: AssignmentBindingType,
-    initialization: Optional[GraphProto],
-    initialization_bindings: Optional[AssignmentBindingType],
+    initialization: GraphProto | None,
+    initialization_bindings: AssignmentBindingType | None,
 ) -> TrainingInfoProto:
     training_info = TrainingInfoProto()
     training_info.algorithm.CopyFrom(algorithm)
@@ -1594,9 +1666,29 @@ def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> int:
     Returns:
         TensorsProto's data_type
     """
-    return cast(
-        int,
-        mapping._NP_TYPE_TO_TENSOR_TYPE[np_dtype],
+    if np_dtype in mapping._NP_TYPE_TO_TENSOR_TYPE:
+        return cast(
+            int,
+            mapping._NP_TYPE_TO_TENSOR_TYPE[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."
     )
 
 
@@ -1609,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:
@@ -1622,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]