onnx 1.16.2__cp312-cp312-win_amd64.whl → 1.18.0__cp312-cp312-win_amd64.whl

onnx/numpy_helper.py

@@ -1,24 +1,34 @@
 # Copyright (c) ONNX Project Contributors
 #
 # SPDX-License-Identifier: Apache-2.0
-
+from __future__ import annotations
 
 import sys
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Any
 
 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
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 
-def combine_pairs_to_complex(fa: Sequence[int]) -> List[complex]:
+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)]
 
 
+@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: Union[np.int16, np.int32, np.ndarray],
-    dims: Optional[Union[int, Sequence[int]]] = None,
+    data: np.int16 | np.int32 | np.ndarray,
+    dims: int | Sequence[int] | None = None,
 ) -> np.ndarray:
     """Converts ndarray of bf16 (as uint32) to f32 (as uint32).
 
@@ -55,6 +65,7 @@ def _float8e4m3_to_float32_scalar(ival: int, fn: bool, uz: bool) -> np.float32:
         if ival == 127:  # noqa: PLR2004
             return np.float32(np.nan)
 
+    ival = np.uint32(ival)  # type: ignore[assignment]
     expo = (ival & 0x78) >> 3
     mant = ival & 0x07
     sign = ival & 0x80
@@ -85,9 +96,13 @@ _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: Union[np.int16, np.int32, np.ndarray],
-    dims: Optional[Union[int, Sequence[int]]] = None,
+    data: np.int16 | np.int32 | np.ndarray,
+    dims: int | Sequence[int] | None = None,
     fn: bool = True,
     uz: bool = False,
 ) -> np.ndarray:
@@ -133,6 +148,7 @@ def _float8e5m2_to_float32_scalar(ival: int, fn: bool, uz: bool) -> np.float32:
     else:
         raise NotImplementedError("fn and uz must be both False or True.")
 
+    ival = np.uint32(ival)  # type: ignore[assignment]
     expo = (ival & 0x7C) >> 2
     mant = ival & 0x03
     sign = ival & 0x80
@@ -159,9 +175,13 @@ _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: Union[np.int16, np.int32, np.ndarray],
-    dims: Optional[Union[int, Sequence[int]]] = None,
+    data: np.int16 | np.int32 | np.ndarray,
+    dims: int | Sequence[int] | None = None,
     fn: bool = False,
     uz: bool = False,
 ) -> np.ndarray:
@@ -185,9 +205,13 @@ def float8e5m2_to_float32(
     return res.reshape(dims)  # type: ignore[no-any-return]
 
 
+@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: Union[np.int32, np.ndarray],
-    dims: Union[int, Sequence[int]],
+    data: np.int32 | np.ndarray,
+    dims: int | Sequence[int],
     signed: bool,
 ) -> np.ndarray:
     """Converts ndarray of int4 (as packed uint8) to f32
@@ -202,7 +226,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())
@@ -218,8 +242,99 @@ def unpack_int4(
     return res
 
 
+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_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(
+    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 _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 _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
     """Converts a tensor def object to a numpy array.
+    Supports types defined in :mod:`onnx._custom_element_types`.
 
     Args:
         tensor: a TensorProto object.
@@ -257,36 +372,43 @@ 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()
 
-        # manually convert bf16 since there's no numpy support
+        data: np.ndarray[Any, Any]
         if tensor_dtype == TensorProto.BFLOAT16:
-            data = np.frombuffer(raw_data, dtype=np.int16)
-            return bfloat16_to_float32(data, dims)
+            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)
-            return float8e4m3_to_float32(data, dims)
+            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)
-            return float8e4m3_to_float32(data, dims, uz=True)
+            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)
-            return float8e5m2_to_float32(data, dims)
+            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)
-            return float8e5m2_to_float32(data, dims, fn=True, uz=True)
+            data = np.frombuffer(raw_data, dtype=np.int8).reshape(dims)
+            return data.view(custom_np_types.float8e5m2fnuz)
 
         if tensor_dtype == TensorProto.UINT4:
             data = np.frombuffer(raw_data, dtype=np.uint8)
-            return unpack_int4(data, dims, signed=False)
+            return _unpack_uint4(
+                data,
+                dims,
+            ).view(custom_np_types.uint4)
 
         if tensor_dtype == TensorProto.INT4:
-            data = np.frombuffer(raw_data, dtype=np.int8)
-            return unpack_int4(data, dims, signed=True)
+            data = np.frombuffer(raw_data, dtype=np.uint8)
+            return _unpack_int4(data, dims).view(custom_np_types.int4)
 
-        return np.frombuffer(raw_data, dtype=np_dtype).reshape(dims)  # type: ignore[no-any-return]
+        if tensor_dtype == TensorProto.FLOAT4E2M1:
+            data = np.frombuffer(raw_data, dtype=np.uint8)
+            return _unpack_uint4(data, dims).view(custom_np_types.float4e2m1)
+
+        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:
@@ -296,43 +418,58 @@ def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PL
             .view(np.float16)
         )
 
-    # 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)
+        data = (
+            np.asarray(tensor.int32_data, dtype=np.int32)
+            .astype(np.uint16)
+            .reshape(dims)
+        )
+        return data.view(custom_np_types.bfloat16)
 
     if tensor_dtype == TensorProto.FLOAT8E4M3FN:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return float8e4m3_to_float32(data, dims)
+        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)
-        return float8e4m3_to_float32(data, dims, uz=True)
+        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)
-        return float8e5m2_to_float32(data, dims)
+        data = (
+            np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8).reshape(dims)
+        )
+        return data.view(custom_np_types.float8e5m2)
 
     if tensor_dtype == TensorProto.FLOAT8E5M2FNUZ:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return float8e5m2_to_float32(data, dims, fn=True, uz=True)
+        data = (
+            np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8).reshape(dims)
+        )
+        return data.view(custom_np_types.float8e5m2fnuz)
 
     if tensor_dtype == TensorProto.UINT4:
-        data = np.asarray(tensor.int32_data, dtype=storage_np_dtype)
-        return unpack_int4(data, dims, signed=False)
+        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=storage_np_dtype)
-        return unpack_int4(data, dims, signed=True)
+        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)
 
     data = getattr(tensor, storage_field)
     if tensor_dtype in (TensorProto.COMPLEX64, TensorProto.COMPLEX128):
-        data = combine_pairs_to_complex(data)  # type: ignore[assignment,arg-type]
+        data = _combine_pairs_to_complex(data)  # type: ignore[assignment,arg-type]
 
     return np.asarray(data, dtype=storage_np_dtype).astype(np_dtype).reshape(dims)
 
 
-def from_array(arr: np.ndarray, name: Optional[str] = None) -> TensorProto:
+def _from_array(arr: np.ndarray, name: str | None = None) -> TensorProto:
     """Converts a numpy array to a tensor def.
 
     Args:
@@ -352,7 +489,7 @@ def from_array(arr: np.ndarray, name: Optional[str] = None) -> TensorProto:
     if name:
         tensor.name = name
 
-    if arr.dtype == object:
+    if arr.dtype == object or np.issubdtype(arr.dtype, np.str_):
         # Special care for strings.
         tensor.data_type = helper.np_dtype_to_tensor_dtype(arr.dtype)
         # TODO: Introduce full string support.
@@ -392,12 +529,76 @@ def from_array(arr: np.ndarray, name: Optional[str] = None) -> TensorProto:
     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)
+        _convert_endian(tensor)
 
     return tensor
 
 
-def to_list(sequence: SequenceProto) -> List[Any]:
+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]:
     """Converts a sequence def to a Python list.
 
     Args:
@@ -408,7 +609,7 @@ def to_list(sequence: SequenceProto) -> List[Any]:
     """
     elem_type = sequence.elem_type
     if elem_type == SequenceProto.TENSOR:
-        return [to_array(v) for v in sequence.tensor_values]  # type: ignore[arg-type]
+        return [to_array(v) for v in sequence.tensor_values]
     if elem_type == SequenceProto.SPARSE_TENSOR:
         return [to_array(v) for v in sequence.sparse_tensor_values]  # type: ignore[arg-type]
     if elem_type == SequenceProto.SEQUENCE:
@@ -419,7 +620,7 @@ def to_list(sequence: SequenceProto) -> List[Any]:
 
 
 def from_list(
-    lst: List[Any], name: Optional[str] = None, dtype: Optional[int] = None
+    lst: list[Any], name: str | None = None, dtype: int | None = None
 ) -> SequenceProto:
     """Converts a list into a sequence def.
 
@@ -475,7 +676,7 @@ def from_list(
     return sequence
 
 
-def to_dict(map_proto: MapProto) -> Dict[Any, Any]:
+def to_dict(map_proto: MapProto) -> dict[Any, Any]:
     """Converts a map def to a Python dictionary.
 
     Args:
@@ -484,7 +685,7 @@ def to_dict(map_proto: MapProto) -> Dict[Any, Any]:
     Returns:
         The converted dictionary.
     """
-    key_list: List[Any] = []
+    key_list: list[Any] = []
     if map_proto.key_type == TensorProto.STRING:
         key_list = list(map_proto.string_keys)
     else:
@@ -501,7 +702,7 @@ def to_dict(map_proto: MapProto) -> Dict[Any, Any]:
     return dictionary
 
 
-def from_dict(dict_: Dict[Any, Any], name: Optional[str] = None) -> MapProto:
+def from_dict(dict_: dict[Any, Any], name: str | None = None) -> MapProto:
     """Converts a Python dictionary into a map def.
 
     Args:
@@ -529,12 +730,7 @@ def from_dict(dict_: Dict[Any, Any], name: Optional[str] = None) -> MapProto:
         TensorProto.UINT64,
     ]
 
-    if not (
-        all(
-            np.result_type(key) == raw_key_type  # type: ignore[arg-type]
-            for key in keys
-        )
-    ):
+    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."
@@ -559,7 +755,7 @@ def from_dict(dict_: Dict[Any, Any], name: Optional[str] = None) -> MapProto:
     return map_proto
 
 
-def to_optional(optional: OptionalProto) -> Optional[Any]:
+def to_optional(optional: OptionalProto) -> Any | None:
     """Converts an optional def to a Python optional.
 
     Args:
@@ -585,7 +781,7 @@ def to_optional(optional: OptionalProto) -> Optional[Any]:
 
 
 def from_optional(
-    opt: Optional[Any], name: Optional[str] = None, dtype: Optional[int] = None
+    opt: Any | None, name: str | None = None, dtype: int | None = None
 ) -> OptionalProto:
     """Converts an optional value into a Optional def.
 
@@ -604,10 +800,9 @@ def from_optional(
     if name:
         optional.name = name
 
-    if dtype:
+    if dtype is not None:
         # dtype must be a valid OptionalProto.DataType
-        valid_dtypes = list(OptionalProto.DataType.values())
-        if dtype not in valid_dtypes:
+        if dtype not in OptionalProto.DataType.values():
             raise TypeError(f"{dtype} must be a valid OptionalProto.DataType.")
         elem_type = dtype
     elif isinstance(opt, dict):
@@ -636,8 +831,8 @@ def from_optional(
     return optional
 
 
-def convert_endian(tensor: TensorProto) -> None:
-    """Call to convert endianess of raw data in tensor.
+def _convert_endian(tensor: TensorProto) -> None:
+    """Call to convert endianness of raw data in tensor.
 
     Args:
         tensor: TensorProto to be converted.
@@ -650,7 +845,7 @@ def convert_endian(tensor: TensorProto) -> None:
 
 
 def create_random_int(
-    input_shape: Tuple[int], dtype: np.dtype, seed: int = 1
+    input_shape: tuple[int], dtype: np.dtype, seed: int = 1
 ) -> np.ndarray:
     """Create random integer array for backend/test/case/node.