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

onnx/numpy_helper.py

@@ -4,19 +4,27 @@
 from __future__ import annotations
 
 import sys
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any
 
+import ml_dtypes
 import numpy as np
+import numpy.typing as npt
+import typing_extensions
 
-import onnx._custom_element_types as custom_np_types
-from onnx import MapProto, OptionalProto, SequenceProto, TensorProto, helper, subbyte
-from onnx.external_data_helper import load_external_data_for_tensor, uses_external_data
+import onnx.external_data_helper
+from onnx import helper, subbyte
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
 
-def combine_pairs_to_complex(fa: Sequence[int]) -> list[complex]:
-    return [complex(fa[i * 2], fa[i * 2 + 1]) for i in range(len(fa) // 2)]
+# System is little endian
+_IS_LITTLE_ENDIAN = sys.byteorder == "little"
 
 
+@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 bfloat16_to_float32(
     data: np.int16 | np.int32 | np.ndarray,
     dims: int | Sequence[int] | None = None,
@@ -87,6 +95,10 @@ _float8e4m3_to_float32 = np.vectorize(
 )
 
 
+@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 float8e4m3_to_float32(
     data: np.int16 | np.int32 | np.ndarray,
     dims: int | Sequence[int] | None = None,
@@ -162,6 +174,10 @@ _float8e5m2_to_float32 = np.vectorize(
 )
 
 
+@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 float8e5m2_to_float32(
     data: np.int16 | np.int32 | np.ndarray,
     dims: int | Sequence[int] | None = None,
@@ -188,6 +204,85 @@ def float8e5m2_to_float32(
     return res.reshape(dims)  # type: ignore[no-any-return]
 
 
+def to_float8e8m0(
+    x: np.ndarray,
+    saturate: bool = True,
+    round_mode: str = "up",
+) -> np.ndarray:
+    """Convert float32 NumPy array to float8e8m0 representation. If the input
+    is not a float32 array, it will be cast to one first.
+
+    Args:
+        x: Input array to convert.
+        saturate: Whether to saturate at max/min float8e8m0 value.
+        round_mode: "nearest", "up", or "down".
+
+    Returns:
+        np.ndarray: Array of ml_dtypes.float8_e8m0fnu values.
+    """
+    x_f32 = np.asarray(x, dtype=np.float32)
+    f_bits = x_f32.view(np.uint32)
+
+    # Extract exponent bits
+    exponent = (f_bits >> 23) & 0xFF
+    exponent = exponent.astype(
+        np.uint16
+    )  # use uint16 to prevent overflow during computation
+
+    # Identify NaN or Inf
+    special_mask = exponent == 0xFF  # noqa: PLR2004
+    output = np.zeros_like(exponent, dtype=np.uint8)
+    output[special_mask] = 0xFF  # Preserve NaN/Inf as max exponent
+
+    # Process normal numbers
+    normal_mask = ~special_mask
+
+    if round_mode == "nearest":
+        # Get guard, round, sticky, and least significant bits
+        g = ((f_bits & 0x400000) > 0).astype(np.uint8)
+        r = ((f_bits & 0x200000) > 0).astype(np.uint8)
+        s = ((f_bits & 0x1FFFFF) > 0).astype(np.uint8)
+        lsb = (exponent > 0).astype(np.uint8)
+
+        round_up = (g == 1) & ((r == 1) | (s == 1) | (lsb == 1))
+
+        increment = np.zeros_like(exponent)
+        increment[round_up & normal_mask] = 1
+
+        if saturate:
+            max_mask = (exponent == 0xFE) & round_up & normal_mask  # noqa: PLR2004
+            increment[max_mask] = 0  # Don't overflow past max value
+
+        exponent += increment
+
+    elif round_mode == "up":
+        has_fraction = (f_bits & 0x4FFFFF) > 0
+        round_up = has_fraction & normal_mask
+
+        if saturate:
+            max_mask = (exponent == 0xFE) & round_up  # noqa: PLR2004
+            round_up[max_mask] = False
+
+        exponent += round_up.astype(np.uint16)
+
+    elif round_mode == "down":
+        pass  # No rounding needed
+
+    else:
+        raise ValueError(f"Unsupported rounding mode: {round_mode}")
+
+    # Clip exponent to uint8 range
+    exponent = exponent.astype(np.uint8)
+
+    output[normal_mask] = exponent[normal_mask]
+
+    return output.view(ml_dtypes.float8_e8m0fnu)
+
+
+@typing_extensions.deprecated(
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider implementing your own unpack logic",
+    category=DeprecationWarning,
+)
 def unpack_int4(
     data: np.int32 | np.ndarray,
     dims: int | Sequence[int],
@@ -221,9 +316,74 @@ def unpack_int4(
     return res
 
 
-def _to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PLR0911
+def _unpacked_float4e2m1_to_float32(
+    x: npt.NDArray[np.uint8],
+) -> npt.NDArray[np.float32]:
+    """Evaluate the numerical value of an array of unpacked float4e2m1 values (as uint8)
+    See :ref:`onnx-detail-int4` for technical details.
+
+    Args:
+        x: an array of uint8 elements representing a float4e2m1 (using the 4 LSB)
+
+    Returns:
+        An array of float32 elements representing the values of the float4e2m1 input.
+    """
+    # x is stored in 4 LSB of int
+    sign = np.where(np.bitwise_and(x, 0x08), -1, 1)
+    mantissa = (x & 0x01).astype(np.float32)
+    exponent = ((x & 0x06) >> 1).astype(np.float32)
+
+    val = np.where(
+        exponent == 0,
+        sign * (mantissa / 2.0),
+        sign * (1.0 + mantissa / 2.0) * 2.0 ** (exponent - 1),
+    )  # denormalized, normalized
+    return val
+
+
+def _unpack_4bit(
+    data: npt.NDArray[np.uint8], dims: Sequence[int]
+) -> npt.NDArray[np.uint8]:
+    """Convert a packed uint4 array to unpacked uint4 array represented as uint8.
+
+    Args:
+        data: A numpy array.
+        dims: The dimensions are used to reshape the unpacked buffer.
+
+    Returns:
+        A numpy array of int8/uint8 reshaped to dims.
+    """
+    result = np.empty([data.size * 2], dtype=data.dtype)
+    array_low = data & np.uint8(0x0F)
+    array_high = data & np.uint8(0xF0)
+    array_high >>= np.uint8(4)
+    result[0::2] = array_low
+    result[1::2] = array_high
+    if result.size == np.prod(dims) + 1:
+        # handle single-element padding due to odd number of elements
+        result = result[:-1]
+    result.resize(dims, refcheck=False)
+    return result
+
+
+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 to_array(tensor: onnx.TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PLR0911
     """Converts a tensor def object to a numpy array.
 
+    This function uses ml_dtypes if the dtype is not a native numpy dtype.
+
     Args:
         tensor: a TensorProto object.
         base_dir: if external tensor exists, base_dir can help to find the path to it
@@ -233,8 +393,8 @@ def _to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: P
     """
     if tensor.HasField("segment"):
         raise ValueError("Currently not supporting loading segments.")
-    if tensor.data_type == TensorProto.UNDEFINED:
-        raise TypeError("The element type in the input tensor is not defined.")
+    if tensor.data_type == onnx.TensorProto.UNDEFINED:
+        raise TypeError("The element type in the input tensor is UNDEFINED.")
 
     tensor_dtype = tensor.data_type
     np_dtype = helper.tensor_dtype_to_np_dtype(tensor_dtype)
@@ -244,14 +404,14 @@ def _to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: P
     storage_field = helper.tensor_dtype_to_field(tensor_dtype)
     dims = tensor.dims
 
-    if tensor.data_type == TensorProto.STRING:
+    if tensor.data_type == onnx.TensorProto.STRING:
         utf8_strings = getattr(tensor, storage_field)
         ss = [s.decode("utf-8") for s in utf8_strings]
         return np.asarray(ss).astype(np_dtype).reshape(dims)
 
     # Load raw data from external tensor if it exists
-    if uses_external_data(tensor):
-        load_external_data_for_tensor(tensor, base_dir)
+    if onnx.external_data_helper.uses_external_data(tensor):
+        onnx.external_data_helper.load_external_data_for_tensor(tensor, base_dir)
 
     if tensor.HasField("raw_data"):
         # Raw_bytes support: using frombuffer.
@@ -260,195 +420,92 @@ def _to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: P
             # Convert endian from little to big
             raw_data = np.frombuffer(raw_data, dtype=np_dtype).byteswap().tobytes()
 
-        # manually convert bf16 since there's no numpy support
-        if tensor_dtype == TensorProto.BFLOAT16:
-            data = np.frombuffer(raw_data, dtype=np.int16)
-            return bfloat16_to_float32(data, dims)
-
-        if tensor_dtype == TensorProto.FLOAT8E4M3FN:
-            data = np.frombuffer(raw_data, dtype=np.int8)
-            return float8e4m3_to_float32(data, dims)
-
-        if tensor_dtype == TensorProto.FLOAT8E4M3FNUZ:
-            data = np.frombuffer(raw_data, dtype=np.int8)
-            return float8e4m3_to_float32(data, dims, uz=True)
-
-        if tensor_dtype == TensorProto.FLOAT8E5M2:
-            data = np.frombuffer(raw_data, dtype=np.int8)
-            return float8e5m2_to_float32(data, dims)
-
-        if tensor_dtype == TensorProto.FLOAT8E5M2FNUZ:
-            data = np.frombuffer(raw_data, dtype=np.int8)
-            return float8e5m2_to_float32(data, dims, fn=True, uz=True)
-
-        if tensor_dtype == TensorProto.UINT4:
+        if tensor_dtype in {
+            onnx.TensorProto.INT4,
+            onnx.TensorProto.UINT4,
+            onnx.TensorProto.FLOAT4E2M1,
+        }:
             data = np.frombuffer(raw_data, dtype=np.uint8)
-            return unpack_int4(data, dims, signed=False)
-
-        if tensor_dtype == TensorProto.INT4:
-            data = np.frombuffer(raw_data, dtype=np.int8)
-            return unpack_int4(data, dims, signed=True)
+            return _unpack_4bit(data, dims).view(np_dtype)
 
-        return np.frombuffer(raw_data, dtype=np_dtype).reshape(dims)  # type: ignore[no-any-return]
+        return np.frombuffer(raw_data, dtype=np_dtype).reshape(dims)
 
-    # float16 is stored as int32 (uint16 type); Need view to get the original value
-    if tensor_dtype == TensorProto.FLOAT16:
+    if tensor_dtype in {
+        onnx.TensorProto.BFLOAT16,
+        onnx.TensorProto.FLOAT16,
+        onnx.TensorProto.INT16,
+        onnx.TensorProto.UINT16,
+    }:
         return (
-            np.asarray(tensor.int32_data, dtype=np.uint16)
+            np.array(tensor.int32_data, dtype=np.int32)
+            .view(np.uint32)
+            .astype(np.uint16)
             .reshape(dims)
-            .view(np.float16)
+            .view(np_dtype)
         )
 
-    # bfloat16 is stored as int32 (uint16 type); no numpy support for bf16
-    if tensor_dtype == TensorProto.BFLOAT16:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return bfloat16_to_float32(data, dims)
-
-    if tensor_dtype == TensorProto.FLOAT8E4M3FN:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return float8e4m3_to_float32(data, dims)
-
-    if tensor_dtype == TensorProto.FLOAT8E4M3FNUZ:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return float8e4m3_to_float32(data, dims, uz=True)
-
-    if tensor_dtype == TensorProto.FLOAT8E5M2:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return float8e5m2_to_float32(data, dims)
-
-    if tensor_dtype == TensorProto.FLOAT8E5M2FNUZ:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return float8e5m2_to_float32(data, dims, fn=True, uz=True)
-
-    if tensor_dtype == TensorProto.UINT4:
-        data = np.asarray(tensor.int32_data, dtype=storage_np_dtype)
-        return unpack_int4(data, dims, signed=False)
+    if tensor_dtype in {
+        onnx.TensorProto.FLOAT8E4M3FN,
+        onnx.TensorProto.FLOAT8E4M3FNUZ,
+        onnx.TensorProto.FLOAT8E5M2,
+        onnx.TensorProto.FLOAT8E5M2FNUZ,
+        onnx.TensorProto.FLOAT8E8M0,
+        onnx.TensorProto.BOOL,
+    }:
+        return (
+            np.array(tensor.int32_data, dtype=np.int32)
+            .view(np.uint32)
+            .astype(np.uint8)
+            .view(np_dtype)
+            .reshape(dims)
+        )
 
-    if tensor_dtype == TensorProto.INT4:
-        data = np.asarray(tensor.int32_data, dtype=storage_np_dtype)
-        return unpack_int4(data, dims, signed=True)
+    if tensor_dtype in {
+        onnx.TensorProto.UINT4,
+        onnx.TensorProto.INT4,
+        onnx.TensorProto.FLOAT4E2M1,
+    }:
+        data = (
+            np.array(tensor.int32_data, dtype=np.int32).view(np.uint32).astype(np.uint8)
+        )
+        return _unpack_4bit(data, dims).view(np_dtype)
 
     data = getattr(tensor, storage_field)
-    if tensor_dtype in (TensorProto.COMPLEX64, TensorProto.COMPLEX128):
-        data = combine_pairs_to_complex(data)  # type: ignore[assignment,arg-type]
+    if tensor_dtype in (onnx.TensorProto.COMPLEX64, onnx.TensorProto.COMPLEX128):
+        return np.array(data, dtype=storage_np_dtype).view(dtype=np_dtype).reshape(dims)
 
     return np.asarray(data, dtype=storage_np_dtype).astype(np_dtype).reshape(dims)
 
 
-def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:
-    """Converts a tensor def object to a numpy array.
-    Supports types defined in :mod:`onnx._custom_element_types`.
-
-    Args:
-        tensor: a TensorProto object.
-        base_dir: if external tensor exists, base_dir can help to find the path to it
-
-    Returns:
-        arr: the converted array.
-    """
-    elem_type = tensor.data_type
-    if elem_type == TensorProto.BFLOAT16:
-        if uses_external_data(tensor):
-            load_external_data_for_tensor(tensor, base_dir)
-        data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        y = np.empty(shape, dtype=custom_np_types.bfloat16).ravel()
-        for i, d in enumerate(data):
-            y[i] = d
-        return y.reshape(shape)
-
-    if elem_type in (
-        TensorProto.FLOAT8E4M3FN,
-        TensorProto.FLOAT8E4M3FNUZ,
-        TensorProto.FLOAT8E5M2,
-        TensorProto.FLOAT8E5M2FNUZ,
-    ):
-        m = {
-            TensorProto.FLOAT8E4M3FN: custom_np_types.float8e4m3fn,
-            TensorProto.FLOAT8E4M3FNUZ: custom_np_types.float8e4m3fnuz,
-            TensorProto.FLOAT8E5M2: custom_np_types.float8e5m2,
-            TensorProto.FLOAT8E5M2FNUZ: custom_np_types.float8e5m2fnuz,
-        }
-
-        if uses_external_data(tensor):
-            load_external_data_for_tensor(tensor, base_dir)
-        if tensor.HasField("raw_data"):
-            data = tensor.raw_data  # type: ignore[assignment]
-        else:
-            data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        y = np.empty(shape, dtype=m[elem_type]).ravel()  # type: ignore[index]
-        for i, d in enumerate(data):
-            y[i] = d
-        return y.reshape(shape)
-
-    if elem_type in (TensorProto.UINT4, TensorProto.INT4):
-        if uses_external_data(tensor):
-            load_external_data_for_tensor(tensor, base_dir)
-        if tensor.HasField("raw_data"):
-            data = tensor.raw_data  # type: ignore[assignment]
-        else:
-            data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        m = {
-            TensorProto.INT4: custom_np_types.int4,
-            TensorProto.UINT4: custom_np_types.uint4,
-        }
-        dtype = m[elem_type]  # type: ignore[index]
-        signed = elem_type == TensorProto.INT4
-        # 2 packed int4 elements must be represented as a single uint8 value.
-        # Therefore, y is np.uint8 (not the dtype to which the int4 maps)
-        y = np.empty(len(data), dtype=np.uint8).ravel()  # type: ignore[assignment]
-        for i, d in enumerate(data):
-            y[i] = d
-        unpacked_data = unpack_int4(y, dims=shape, signed=signed)
-        return unpacked_data.astype(dtype)
-
-    return _to_array(tensor, base_dir=base_dir)
-
-
-def _from_array(arr: np.ndarray, name: str | None = None) -> TensorProto:
-    """Converts a numpy array to a tensor def.
+def from_array(array: np.ndarray, /, name: str | None = None) -> onnx.TensorProto:
+    """Converts an array into a TensorProto including
 
     Args:
-        arr: a numpy array.
+        array: a numpy array.
         name: (optional) the name of the tensor.
 
     Returns:
         TensorProto: the converted tensor def.
     """
-    if not isinstance(arr, (np.ndarray, np.generic)):
-        raise TypeError(
-            f"arr must be of type np.generic or np.ndarray, got {type(arr)}"
-        )
-
-    tensor = TensorProto()
-    tensor.dims.extend(arr.shape)
+    tensor = onnx.TensorProto()
+    tensor.dims.extend(array.shape)
     if name:
         tensor.name = name
-
-    if arr.dtype == object or np.issubdtype(arr.dtype, np.str_):
+    if array.dtype == object or np.issubdtype(array.dtype, np.str_):
         # Special care for strings.
-        tensor.data_type = helper.np_dtype_to_tensor_dtype(arr.dtype)
+        tensor.data_type = onnx.TensorProto.STRING
         # TODO: Introduce full string support.
-        # We flatten the array in case there are 2-D arrays are specified
-        # We throw the error below if we have a 3-D array or some kind of other
-        # object. If you want more complex shapes then follow the below instructions.
+        # We flatten the array in case there are n-D arrays are specified
+        # If you want more complex shapes then follow the below instructions.
         # Unlike other types where the shape is automatically inferred from
         # nested arrays of values, the only reliable way now to feed strings
         # is to put them into a flat array then specify type astype(object)
         # (otherwise all strings may have different types depending on their length)
         # and then specify shape .reshape([x, y, z])
-        flat_array = arr.flatten()
+        flat_array = array.flatten()
         for e in flat_array:
             if isinstance(e, str):
                 tensor.string_data.append(e.encode("utf-8"))
-            elif isinstance(e, np.ndarray):
-                for s in e:
-                    if isinstance(s, str):
-                        tensor.string_data.append(s.encode("utf-8"))
-                    elif isinstance(s, bytes):
-                        tensor.string_data.append(s)
             elif isinstance(e, bytes):
                 tensor.string_data.append(e)
             else:
@@ -458,82 +515,23 @@ def _from_array(arr: np.ndarray, name: str | None = None) -> TensorProto:
                 )
         return tensor
 
-    # For numerical types, directly use numpy raw bytes.
-    try:
-        dtype = helper.np_dtype_to_tensor_dtype(arr.dtype)
-    except KeyError as e:
-        raise RuntimeError(f"Numpy data type not understood yet: {arr.dtype!r}") from e
+    dtype = helper.np_dtype_to_tensor_dtype(array.dtype)
+    if dtype in {
+        onnx.TensorProto.INT4,
+        onnx.TensorProto.UINT4,
+        onnx.TensorProto.FLOAT4E2M1,
+    }:
+        # Pack the array into int4
+        array = _pack_4bitx2(array)
+    if not _IS_LITTLE_ENDIAN:
+        array = array.view(array.dtype.newbyteorder("<"))
+
+    tensor.raw_data = array.tobytes()
     tensor.data_type = dtype
-    tensor.raw_data = arr.tobytes()  # note: tobytes() is only after 1.9.
-    if sys.byteorder == "big":
-        # Convert endian from big to little
-        convert_endian(tensor)
-
     return tensor
 
 
-def from_array(tensor: np.ndarray, name: str | None = None) -> TensorProto:
-    """Converts an array into a TensorProto including
-    supported type defined in :mod:`onnx._custom_element_types`.
-
-    Args:
-        tensor: a numpy array.
-        name: (optional) the name of the tensor.
-
-    Returns:
-        TensorProto: the converted tensor def.
-    """
-    if not isinstance(tensor, np.ndarray):
-        return _from_array(tensor, name)
-
-    dt = tensor.dtype
-    if dt == custom_np_types.float8e4m3fn and dt.descr[0][0] == "e4m3fn":
-        to = TensorProto.FLOAT8E4M3FN
-        dt_to = np.uint8  # type: ignore[assignment]
-    elif dt == custom_np_types.float8e4m3fnuz and dt.descr[0][0] == "e4m3fnuz":
-        to = TensorProto.FLOAT8E4M3FNUZ
-        dt_to = np.uint8  # type: ignore[assignment]
-    elif dt == custom_np_types.float8e5m2 and dt.descr[0][0] == "e5m2":
-        to = TensorProto.FLOAT8E5M2
-        dt_to = np.uint8  # type: ignore[assignment]
-    elif dt == custom_np_types.float8e5m2fnuz and dt.descr[0][0] == "e5m2fnuz":
-        to = TensorProto.FLOAT8E5M2FNUZ
-        dt_to = np.uint8  # type: ignore[assignment]
-    elif dt == custom_np_types.bfloat16 and dt.descr[0][0] == "bfloat16":
-        to = TensorProto.BFLOAT16
-        dt_to = np.uint16  # type: ignore[assignment]
-    elif dt == custom_np_types.int4 and dt.descr[0][0] == "int4":
-        to = TensorProto.INT4
-        dt_to = np.int8  # type: ignore[assignment]
-    elif dt == custom_np_types.uint4 and dt.descr[0][0] == "uint4":
-        to = TensorProto.UINT4
-        dt_to = np.uint8  # type: ignore[assignment]
-    else:
-        return _from_array(tensor, name)
-
-    if to in (TensorProto.UINT4, TensorProto.INT4):
-        value = tensor.astype(dt_to).ravel()
-        if value.size % 2 == 1:
-            raise ValueError(
-                f"The conversion of a tensor of INT4 or UINT4 requires an even size "
-                f"(shape={tensor.shape}). Every byte stores two INT4 or two UINT4."
-            )
-        buffer = (value[::2] & 0x0F) + (value[1::2] << 4)
-
-        pb = TensorProto()
-        pb.dims.extend(tensor.shape)
-        if name:
-            pb.name = name
-        pb.raw_data = buffer.tobytes()
-        pb.data_type = to
-        return pb
-
-    t = _from_array(tensor.astype(dt_to), name)
-    t.data_type = to
-    return t
-
-
-def to_list(sequence: SequenceProto) -> list[Any]:
+def to_list(sequence: onnx.SequenceProto) -> list[Any]:
     """Converts a sequence def to a Python list.
 
     Args:
@@ -543,20 +541,20 @@ def to_list(sequence: SequenceProto) -> list[Any]:
         list: the converted list.
     """
     elem_type = sequence.elem_type
-    if elem_type == SequenceProto.TENSOR:
-        return [to_array(v) for v in sequence.tensor_values]  # type: ignore[arg-type]
-    if elem_type == SequenceProto.SPARSE_TENSOR:
+    if elem_type == onnx.SequenceProto.TENSOR:
+        return [to_array(v) for v in sequence.tensor_values]
+    if elem_type == onnx.SequenceProto.SPARSE_TENSOR:
         return [to_array(v) for v in sequence.sparse_tensor_values]  # type: ignore[arg-type]
-    if elem_type == SequenceProto.SEQUENCE:
+    if elem_type == onnx.SequenceProto.SEQUENCE:
         return [to_list(v) for v in sequence.sequence_values]
-    if elem_type == SequenceProto.MAP:
+    if elem_type == onnx.SequenceProto.MAP:
         return [to_dict(v) for v in sequence.map_values]
     raise TypeError("The element type in the input sequence is not supported.")
 
 
 def from_list(
     lst: list[Any], name: str | None = None, dtype: int | None = None
-) -> SequenceProto:
+) -> onnx.SequenceProto:
     """Converts a list into a sequence def.
 
     Args:
@@ -568,7 +566,7 @@ def from_list(
     Returns:
         SequenceProto: the converted sequence def.
     """
-    sequence = SequenceProto()
+    sequence = onnx.SequenceProto()
     if name:
         sequence.name = name
 
@@ -577,15 +575,15 @@ def from_list(
     elif len(lst) > 0:
         first_elem = lst[0]
         if isinstance(first_elem, dict):
-            elem_type = SequenceProto.MAP
+            elem_type = onnx.SequenceProto.MAP
         elif isinstance(first_elem, list):
-            elem_type = SequenceProto.SEQUENCE
+            elem_type = onnx.SequenceProto.SEQUENCE
         else:
-            elem_type = SequenceProto.TENSOR
+            elem_type = onnx.SequenceProto.TENSOR
     else:
         # if empty input list and no dtype specified
         # choose sequence of tensors on default
-        elem_type = SequenceProto.TENSOR
+        elem_type = onnx.SequenceProto.TENSOR
     sequence.elem_type = elem_type
 
     if (len(lst) > 0) and not all(isinstance(elem, type(lst[0])) for elem in lst):
@@ -594,13 +592,13 @@ def from_list(
             "for all elements and therefore is not supported as a sequence."
         )
 
-    if elem_type == SequenceProto.TENSOR:
+    if elem_type == onnx.SequenceProto.TENSOR:
         for tensor in lst:
-            sequence.tensor_values.extend([from_array(tensor)])
-    elif elem_type == SequenceProto.SEQUENCE:
+            sequence.tensor_values.extend([from_array(np.asarray(tensor))])
+    elif elem_type == onnx.SequenceProto.SEQUENCE:
         for seq in lst:
             sequence.sequence_values.extend([from_list(seq)])
-    elif elem_type == SequenceProto.MAP:
+    elif elem_type == onnx.SequenceProto.MAP:
         for mapping in lst:
             sequence.map_values.extend([from_dict(mapping)])
     else:
@@ -611,7 +609,7 @@ def from_list(
     return sequence
 
 
-def to_dict(map_proto: MapProto) -> dict[Any, Any]:
+def to_dict(map_proto: onnx.MapProto) -> dict[Any, Any]:
     """Converts a map def to a Python dictionary.
 
     Args:
@@ -621,7 +619,7 @@ def to_dict(map_proto: MapProto) -> dict[Any, Any]:
         The converted dictionary.
     """
     key_list: list[Any] = []
-    if map_proto.key_type == TensorProto.STRING:
+    if map_proto.key_type == onnx.TensorProto.STRING:
         key_list = list(map_proto.string_keys)
     else:
         key_list = list(map_proto.keys)
@@ -637,7 +635,7 @@ def to_dict(map_proto: MapProto) -> dict[Any, Any]:
     return dictionary
 
 
-def from_dict(dict_: dict[Any, Any], name: str | None = None) -> MapProto:
+def from_dict(dict_: dict[Any, Any], name: str | None = None) -> onnx.MapProto:
     """Converts a Python dictionary into a map def.
 
     Args:
@@ -647,30 +645,25 @@ def from_dict(dict_: dict[Any, Any], name: str | None = None) -> MapProto:
     Returns:
         MapProto: the converted map def.
     """
-    map_proto = MapProto()
+    map_proto = onnx.MapProto()
     if name:
         map_proto.name = name
     keys = list(dict_)
     raw_key_type = np.result_type(keys[0])
     key_type = helper.np_dtype_to_tensor_dtype(raw_key_type)
 
-    valid_key_int_types = [
-        TensorProto.INT8,
-        TensorProto.INT16,
-        TensorProto.INT32,
-        TensorProto.INT64,
-        TensorProto.UINT8,
-        TensorProto.UINT16,
-        TensorProto.UINT32,
-        TensorProto.UINT64,
-    ]
-
-    if not (
-        all(
-            np.result_type(key) == raw_key_type  # type: ignore[arg-type]
-            for key in keys
-        )
-    ):
+    valid_key_int_types = {
+        onnx.TensorProto.INT8,
+        onnx.TensorProto.INT16,
+        onnx.TensorProto.INT32,
+        onnx.TensorProto.INT64,
+        onnx.TensorProto.UINT8,
+        onnx.TensorProto.UINT16,
+        onnx.TensorProto.UINT32,
+        onnx.TensorProto.UINT64,
+    }
+
+    if not (all(np.result_type(key) == raw_key_type for key in keys)):
         raise TypeError(
             "The key type in the input dictionary is not the same "
             "for all keys and therefore is not valid as a map."
@@ -687,7 +680,7 @@ def from_dict(dict_: dict[Any, Any], name: str | None = None) -> MapProto:
     value_seq = from_list(values)
 
     map_proto.key_type = key_type
-    if key_type == TensorProto.STRING:
+    if key_type == onnx.TensorProto.STRING:
         map_proto.string_keys.extend(keys)
     elif key_type in valid_key_int_types:
         map_proto.keys.extend(keys)
@@ -695,7 +688,7 @@ def from_dict(dict_: dict[Any, Any], name: str | None = None) -> MapProto:
     return map_proto
 
 
-def to_optional(optional: OptionalProto) -> Any | None:
+def to_optional(optional: onnx.OptionalProto) -> Any | None:
     """Converts an optional def to a Python optional.
 
     Args:
@@ -705,24 +698,24 @@ def to_optional(optional: OptionalProto) -> Any | None:
         opt: the converted optional.
     """
     elem_type = optional.elem_type
-    if elem_type == OptionalProto.UNDEFINED:
+    if elem_type == onnx.OptionalProto.UNDEFINED:
         return None
-    if elem_type == OptionalProto.TENSOR:
+    if elem_type == onnx.OptionalProto.TENSOR:
         return to_array(optional.tensor_value)
-    if elem_type == OptionalProto.SPARSE_TENSOR:
+    if elem_type == onnx.OptionalProto.SPARSE_TENSOR:
         return to_array(optional.sparse_tensor_value)  # type: ignore[arg-type]
-    if elem_type == OptionalProto.SEQUENCE:
+    if elem_type == onnx.OptionalProto.SEQUENCE:
         return to_list(optional.sequence_value)
-    if elem_type == OptionalProto.MAP:
+    if elem_type == onnx.OptionalProto.MAP:
         return to_dict(optional.map_value)
-    if elem_type == OptionalProto.OPTIONAL:
+    if elem_type == onnx.OptionalProto.OPTIONAL:
         return to_optional(optional.optional_value)
     raise TypeError("The element type in the input optional is not supported.")
 
 
 def from_optional(
     opt: Any | None, name: str | None = None, dtype: int | None = None
-) -> OptionalProto:
+) -> onnx.OptionalProto:
     """Converts an optional value into a Optional def.
 
     Args:
@@ -736,33 +729,32 @@ def from_optional(
         optional: the converted optional def.
     """
     # TODO: create a map and replace conditional branches
-    optional = OptionalProto()
+    optional = onnx.OptionalProto()
     if name:
         optional.name = name
 
-    if dtype:
-        # dtype must be a valid OptionalProto.DataType
-        valid_dtypes = list(OptionalProto.DataType.values())
-        if dtype not in valid_dtypes:
+    if dtype is not None:
+        # dtype must be a valid onnx.OptionalProto.DataType
+        if dtype not in onnx.OptionalProto.DataType.values():
             raise TypeError(f"{dtype} must be a valid OptionalProto.DataType.")
         elem_type = dtype
     elif isinstance(opt, dict):
-        elem_type = OptionalProto.MAP
+        elem_type = onnx.OptionalProto.MAP
     elif isinstance(opt, list):
-        elem_type = OptionalProto.SEQUENCE
+        elem_type = onnx.OptionalProto.SEQUENCE
     elif opt is None:
-        elem_type = OptionalProto.UNDEFINED
+        elem_type = onnx.OptionalProto.UNDEFINED
     else:
-        elem_type = OptionalProto.TENSOR
+        elem_type = onnx.OptionalProto.TENSOR
 
     optional.elem_type = elem_type
 
     if opt is not None:
-        if elem_type == OptionalProto.TENSOR:
+        if elem_type == onnx.OptionalProto.TENSOR:
             optional.tensor_value.CopyFrom(from_array(opt))
-        elif elem_type == OptionalProto.SEQUENCE:
+        elif elem_type == onnx.OptionalProto.SEQUENCE:
             optional.sequence_value.CopyFrom(from_list(opt))
-        elif elem_type == OptionalProto.MAP:
+        elif elem_type == onnx.OptionalProto.MAP:
             optional.map_value.CopyFrom(from_dict(opt))
         else:
             raise TypeError(
@@ -772,19 +764,6 @@ def from_optional(
     return optional
 
 
-def convert_endian(tensor: TensorProto) -> None:
-    """Call to convert endianness of raw data in tensor.
-
-    Args:
-        tensor: TensorProto to be converted.
-    """
-    tensor_dtype = tensor.data_type
-    np_dtype = helper.tensor_dtype_to_np_dtype(tensor_dtype)
-    tensor.raw_data = (
-        np.frombuffer(tensor.raw_data, dtype=np_dtype).byteswap().tobytes()
-    )
-
-
 def create_random_int(
     input_shape: tuple[int], dtype: np.dtype, seed: int = 1
 ) -> np.ndarray:
@@ -815,3 +794,19 @@ def create_random_int(
         return np.random.randint(start, end, size=input_shape).astype(dtype)
     else:
         raise TypeError(f"{dtype} is not supported by create_random_int.")
+
+
+def saturate_cast(x: np.ndarray, dtype: np.dtype) -> np.ndarray:
+    """Saturate cast for numeric types.
+
+    This function ensures that values outside the representable range
+    of the target dtype are clamped to the maximum or minimum representable
+    value of that dtype.
+    """
+    if np.issubdtype(dtype, np.integer) or dtype in (ml_dtypes.int4, ml_dtypes.uint4):
+        info = ml_dtypes.iinfo(dtype)
+        x = np.round(x)
+    else:
+        info = ml_dtypes.finfo(dtype)  # type: ignore[assignment]
+
+    return np.clip(x, info.min, info.max).astype(dtype)