onnx 1.18.0__cp313-cp313-macosx_12_0_universal2.whl → 1.19.1rc1__cp313-cp313-macosx_12_0_universal2.whl

onnx/numpy_helper.py

@@ -6,20 +6,19 @@ from __future__ import annotations
 import sys
 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(
@@ -205,6 +204,81 @@ 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,
@@ -226,7 +300,7 @@ def unpack_int4(
     Returns:
         A numpy array of float32 reshaped to dims.
     """
-    single_func = lambda x: subbyte._unpack_single_4bitx2(x, signed)  # noqa: E731
+    single_func = lambda x: subbyte.unpack_single_4bitx2(x, signed)  # noqa: E731
     func = np.frompyfunc(single_func, 1, 2)
 
     res_high, res_low = func(data.ravel())
@@ -242,7 +316,9 @@ def unpack_int4(
     return res
 
 
-def unpacked_float4e2m1_to_float32(x: npt.NDArray[np.uint8]) -> npt.NDArray[np.float32]:
+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.
 
@@ -265,26 +341,7 @@ def unpacked_float4e2m1_to_float32(x: npt.NDArray[np.uint8]) -> npt.NDArray[np.f
     return val
 
 
-def unpack_float4e2m1(
-    data: npt.NDArray[np.uint8],
-    dims: Sequence[int],
-) -> np.ndarray:
-    """Converts ndarray of float4e2m1 (as packed uint8) to f32
-    See :ref:`onnx-detail-float4` for technical details.
-
-    Args:
-        data: A numpy array, empty dimensions are allowed if dims is
-            None.
-        dims: The dimensions are used to reshape the unpacked buffer
-
-    Returns:
-        A numpy array of float32 reshaped to dims.
-    """
-    res = _unpack_uint4(data, dims)
-    return unpacked_float4e2m1_to_float32(res)
-
-
-def _unpack_uint4(
+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.
@@ -309,32 +366,23 @@ def _unpack_uint4(
     return result
 
 
-def _extend_int4_sign_bits(x: npt.NDArray[np.uint8]) -> npt.NDArray[np.int8]:
-    """Extend 4-bit signed integer to 8-bit signed integer."""
-    return np.where((x >> 3) == 0, x, x | 0xF0).astype(np.int8)
+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 _unpack_int4(
-    data: npt.NDArray[np.uint8], dims: Sequence[int]
-) -> npt.NDArray[np.int8]:
-    """Convert a packed (signed) int4 array to unpacked int4 array represented as int8.
-
-    The sign bit is extended to the most significant bit of the int8.
-
-    Args:
-        data: A numpy array.
-        dims: The dimensions are used to reshape the unpacked buffer.
-
-    Returns:
-        A numpy array of int8 reshaped to dims.
-    """
-    unpacked = _unpack_uint4(data, dims)
-    return _extend_int4_sign_bits(unpacked)
-
-
-def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PLR0911
+def to_array(tensor: onnx.TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PLR0911
     """Converts a tensor def object to a numpy array.
-    Supports types defined in :mod:`onnx._custom_element_types`.
+
+    This function uses ml_dtypes if the dtype is not a native numpy dtype.
 
     Args:
         tensor: a TensorProto object.
@@ -345,8 +393,8 @@ def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PL
     """
     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)
@@ -356,14 +404,14 @@ def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PL
     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.
@@ -372,145 +420,92 @@ def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PL
             # Convert endian from little to big
             raw_data = np.frombuffer(raw_data, dtype=np_dtype).byteswap().tobytes()
 
-        data: np.ndarray[Any, Any]
-        if tensor_dtype == TensorProto.BFLOAT16:
-            data = np.frombuffer(raw_data, dtype=np.int16).reshape(dims)
-            return data.view(custom_np_types.bfloat16)
-
-        if tensor_dtype == TensorProto.FLOAT8E4M3FN:
-            data = np.frombuffer(raw_data, dtype=np.int8).reshape(dims)
-            return data.view(custom_np_types.float8e4m3fn)
-
-        if tensor_dtype == TensorProto.FLOAT8E4M3FNUZ:
-            data = np.frombuffer(raw_data, dtype=np.int8).reshape(dims)
-            return data.view(custom_np_types.float8e4m3fnuz)
-
-        if tensor_dtype == TensorProto.FLOAT8E5M2:
-            data = np.frombuffer(raw_data, dtype=np.int8).reshape(dims)
-            return data.view(custom_np_types.float8e5m2)
-
-        if tensor_dtype == TensorProto.FLOAT8E5M2FNUZ:
-            data = np.frombuffer(raw_data, dtype=np.int8).reshape(dims)
-            return data.view(custom_np_types.float8e5m2fnuz)
-
-        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_uint4(
-                data,
-                dims,
-            ).view(custom_np_types.uint4)
-
-        if tensor_dtype == TensorProto.INT4:
-            data = np.frombuffer(raw_data, dtype=np.uint8)
-            return _unpack_int4(data, dims).view(custom_np_types.int4)
-
-        if tensor_dtype == TensorProto.FLOAT4E2M1:
-            data = np.frombuffer(raw_data, dtype=np.uint8)
-            return _unpack_uint4(data, dims).view(custom_np_types.float4e2m1)
+            return _unpack_4bit(data, dims).view(np_dtype)
 
         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)
-            .reshape(dims)
-            .view(np.float16)
-        )
-
-    if tensor_dtype == TensorProto.BFLOAT16:
-        data = (
-            np.asarray(tensor.int32_data, dtype=np.int32)
+            np.array(tensor.int32_data, dtype=np.int32)
+            .view(np.uint32)
             .astype(np.uint16)
             .reshape(dims)
+            .view(np_dtype)
         )
-        return data.view(custom_np_types.bfloat16)
-
-    if tensor_dtype == TensorProto.FLOAT8E4M3FN:
-        data = (
-            np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8).reshape(dims)
-        )
-        return data.view(custom_np_types.float8e4m3fn)
-
-    if tensor_dtype == TensorProto.FLOAT8E4M3FNUZ:
-        data = (
-            np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8).reshape(dims)
-        )
-        return data.view(custom_np_types.float8e4m3fnuz)
 
-    if tensor_dtype == TensorProto.FLOAT8E5M2:
-        data = (
-            np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8).reshape(dims)
+    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)
         )
-        return data.view(custom_np_types.float8e5m2)
 
-    if tensor_dtype == TensorProto.FLOAT8E5M2FNUZ:
+    if tensor_dtype in {
+        onnx.TensorProto.UINT4,
+        onnx.TensorProto.INT4,
+        onnx.TensorProto.FLOAT4E2M1,
+    }:
         data = (
-            np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8).reshape(dims)
+            np.array(tensor.int32_data, dtype=np.int32).view(np.uint32).astype(np.uint8)
         )
-        return data.view(custom_np_types.float8e5m2fnuz)
-
-    if tensor_dtype == TensorProto.UINT4:
-        data = np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8)
-        return _unpack_uint4(data, dims).view(custom_np_types.uint4)
-
-    if tensor_dtype == TensorProto.INT4:
-        data = np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8)
-        return _unpack_int4(data, dims).view(custom_np_types.int4)
-
-    if tensor_dtype == TensorProto.FLOAT4E2M1:
-        data = np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8)
-        return _unpack_uint4(data, dims).view(custom_np_types.float4e2m1)
+        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 _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:
@@ -520,85 +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
-    elif dt == custom_np_types.float8e4m3fnuz and dt.descr[0][0] == "e4m3fnuz":
-        to = TensorProto.FLOAT8E4M3FNUZ
-        dt_to = np.uint8
-    elif dt == custom_np_types.float8e5m2 and dt.descr[0][0] == "e5m2":
-        to = TensorProto.FLOAT8E5M2
-        dt_to = np.uint8
-    elif dt == custom_np_types.float8e5m2fnuz and dt.descr[0][0] == "e5m2fnuz":
-        to = TensorProto.FLOAT8E5M2FNUZ
-        dt_to = np.uint8
-    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
-    elif dt == custom_np_types.float4e2m1 and dt.descr[0][0] == "float4e2m1":
-        to = TensorProto.FLOAT4E2M1
-        dt_to = np.uint8
-    else:
-        return _from_array(tensor, name)
-
-    if to in (TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1):
-        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:
@@ -608,20 +541,20 @@ def to_list(sequence: SequenceProto) -> list[Any]:
         list: the converted list.
     """
     elem_type = sequence.elem_type
-    if elem_type == SequenceProto.TENSOR:
+    if elem_type == onnx.SequenceProto.TENSOR:
         return [to_array(v) for v in sequence.tensor_values]
-    if elem_type == SequenceProto.SPARSE_TENSOR:
+    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:
@@ -633,7 +566,7 @@ def from_list(
     Returns:
         SequenceProto: the converted sequence def.
     """
-    sequence = SequenceProto()
+    sequence = onnx.SequenceProto()
     if name:
         sequence.name = name
 
@@ -642,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):
@@ -659,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:
@@ -676,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:
@@ -686,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)
@@ -702,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:
@@ -712,23 +645,23 @@ 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,
-    ]
+    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(
@@ -747,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)
@@ -755,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:
@@ -765,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:
@@ -796,32 +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 is not None:
-        # dtype must be a valid OptionalProto.DataType
-        if dtype not in OptionalProto.DataType.values():
+        # 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(
@@ -831,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:
@@ -874,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)