onnx 1.17.0__cp311-cp311-macosx_12_0_universal2.whl → 1.19.0__cp311-cp311-macosx_12_0_universal2.whl

onnx/helper.py

@@ -4,51 +4,44 @@
 from __future__ import annotations
 
 import collections.abc
+import functools
+import math
 import numbers
 import struct
+import typing
 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
 
 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,
 )
 
-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,12 +71,14 @@ 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),
+    ("1.19.0", 12, 24, 5, 1),
 ]
 
-VersionMapType = Dict[Tuple[str, int], int]
+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 = {}
 
@@ -100,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(
@@ -206,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
 
@@ -218,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
     """
@@ -304,11 +299,12 @@ 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 = 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 +312,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 +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)
-    ]
-
-
-# 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.
+@typing_extensions.deprecated(
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion"
+)
 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
@@ -386,6 +377,9 @@ def float32_to_bfloat16(fval: float, truncate: bool = False) -> int:
     return (ival + rounded) >> 16
 
 
+@typing_extensions.deprecated(
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion"
+)
 def float32_to_float8e4m3(  # noqa: PLR0911
     fval: float,
     scale: float = 1.0,
@@ -520,6 +514,9 @@ def float32_to_float8e4m3(  # noqa: PLR0911
         return int(ret)
 
 
+@typing_extensions.deprecated(
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion"
+)
 def float32_to_float8e5m2(  # noqa: PLR0911
     fval: float,
     scale: float = 1.0,
@@ -646,6 +643,9 @@ 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"
+)
 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,20 +660,64 @@ def pack_float32_to_4bit(array: np.ndarray | Sequence, signed: bool) -> np.ndarr
     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]))
 
-    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"
+)
+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: 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])
+    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
@@ -682,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
@@ -695,106 +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 = 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. 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.
+        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 type(vals) is 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)
-            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.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:
@@ -805,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)
@@ -824,16 +847,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 +900,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 +909,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(
@@ -932,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):
@@ -955,7 +984,7 @@ 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),
             ):
@@ -1354,13 +1383,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)  # 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)))
             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')}"
@@ -1375,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)))
@@ -1420,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
@@ -1438,9 +1473,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))
@@ -1550,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:
@@ -1562,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:
@@ -1574,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.
 
@@ -1586,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:
@@ -1600,25 +1644,13 @@ 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
-
-    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,
-    }:
-        return custom_np_types.mapping_name_to_data_type[np_dtype.descr[0][0]]
+        return TensorProto.STRING  # type: ignore[no-any-return]
 
     raise ValueError(
         f"Unable to convert type {np_dtype!r} into TensorProto element type."
@@ -1631,10 +1663,13 @@ 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 = {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()  # type: ignore[attr-defined]
+}
 
 
 def _attr_type_to_str(attr_type: int) -> str:
@@ -1646,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():
-        return _ATTRIBUTE_TYPE_TO_STR[attr_type]  # type: ignore[no-any-return]
-    return AttributeProto.AttributeType.keys()[0]  # type: ignore[no-any-return]
+    if attr_type in AttributeProto.AttributeType.values():  # type: ignore[attr-defined]
+        return _ATTRIBUTE_TYPE_TO_STR[attr_type]
+    return AttributeProto.AttributeType.keys()[0]  # type: ignore