onnx-ir 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

onnx_ir/__init__.py

@@ -167,4 +167,4 @@ def __set_module() -> None:
 
 
 __set_module()
-__version__ = "0.1.3"
+__version__ = "0.1.4"

onnx_ir/_core.py

@@ -657,15 +657,13 @@ class ExternalTensor(TensorBase, _protocols.TensorProtocol):  # pylint: disable=
             self._array = np.empty(self.shape.numpy(), dtype=self.dtype.numpy())
             return
         # Map the whole file into the memory
-        # TODO(justinchuby): Verify if this would exhaust the memory address space
         with open(self.path, "rb") as f:
             self.raw = mmap.mmap(
                 f.fileno(),
                 0,
                 access=mmap.ACCESS_READ,
             )
-        # Handle the byte order correctly by always using little endian
-        dt = np.dtype(self.dtype.numpy()).newbyteorder("<")
+
         if self.dtype in {
             _enums.DataType.INT4,
             _enums.DataType.UINT4,
@@ -675,16 +673,18 @@ class ExternalTensor(TensorBase, _protocols.TensorProtocol):  # pylint: disable=
             dt = np.dtype(np.uint8).newbyteorder("<")
             count = self.size // 2 + self.size % 2
         else:
+            # Handle the byte order correctly by always using little endian
+            dt = np.dtype(self.dtype.numpy()).newbyteorder("<")
             count = self.size
+
         self._array = np.frombuffer(self.raw, dtype=dt, offset=self.offset or 0, count=count)
         shape = self.shape.numpy()
-        if self.dtype == _enums.DataType.INT4:
-            # Unpack the int4 arrays
-            self._array = _type_casting.unpack_int4(self._array, shape)
-        elif self.dtype == _enums.DataType.UINT4:
-            self._array = _type_casting.unpack_uint4(self._array, shape)
-        elif self.dtype == _enums.DataType.FLOAT4E2M1:
-            self._array = _type_casting.unpack_float4e2m1(self._array, shape)
+
+        if self.dtype.bitwidth == 4:
+            # Unpack the 4bit arrays
+            self._array = _type_casting.unpack_4bitx2(self._array, shape).view(
+                self.dtype.numpy()
+            )
         else:
             self._array = self._array.reshape(shape)
 
@@ -1071,15 +1071,7 @@ class PackedTensor(TensorBase, _protocols.TensorProtocol, Generic[TArrayCompatib
         """
         array = self.numpy_packed()
         # ONNX IR returns the unpacked arrays
-        if self.dtype == _enums.DataType.INT4:
-            return _type_casting.unpack_int4(array, self.shape.numpy())
-        if self.dtype == _enums.DataType.UINT4:
-            return _type_casting.unpack_uint4(array, self.shape.numpy())
-        if self.dtype == _enums.DataType.FLOAT4E2M1:
-            return _type_casting.unpack_float4e2m1(array, self.shape.numpy())
-        raise TypeError(
-            f"PackedTensor only supports INT4, UINT4, FLOAT4E2M1, but got {self.dtype}"
-        )
+        return _type_casting.unpack_4bitx2(array, self.shape.numpy()).view(self.dtype.numpy())
 
     def numpy_packed(self) -> npt.NDArray[np.uint8]:
         """Return the tensor as a packed array."""

onnx_ir/_enums.py

@@ -169,6 +169,48 @@ class DataType(enum.IntEnum):
             DataType.FLOAT4E2M1,
         }
 
+    def is_integer(self) -> bool:
+        """Returns True if the data type is an integer.
+
+        .. versionadded:: 0.1.4
+        """
+        return self in {
+            DataType.UINT8,
+            DataType.INT8,
+            DataType.UINT16,
+            DataType.INT16,
+            DataType.INT32,
+            DataType.INT64,
+            DataType.UINT32,
+            DataType.UINT64,
+            DataType.UINT4,
+            DataType.INT4,
+        }
+
+    def is_signed(self) -> bool:
+        """Returns True if the data type is a signed type.
+
+        .. versionadded:: 0.1.4
+        """
+        return self in {
+            DataType.FLOAT,
+            DataType.INT8,
+            DataType.INT16,
+            DataType.INT32,
+            DataType.INT64,
+            DataType.FLOAT16,
+            DataType.DOUBLE,
+            DataType.COMPLEX64,
+            DataType.COMPLEX128,
+            DataType.BFLOAT16,
+            DataType.FLOAT8E4M3FN,
+            DataType.FLOAT8E4M3FNUZ,
+            DataType.FLOAT8E5M2,
+            DataType.FLOAT8E5M2FNUZ,
+            DataType.INT4,
+            DataType.FLOAT4E2M1,
+        }
+
     def __repr__(self) -> str:
         return self.name
 

onnx_ir/_type_casting.py

@@ -1,14 +1,12 @@
 # Copyright (c) ONNX Project Contributors
 # SPDX-License-Identifier: Apache-2.0
 """Numpy utilities for non-native type operation."""
-# TODO(justinchuby): Upstream the logic to onnx
 
 from __future__ import annotations
 
 import typing
 from collections.abc import Sequence
 
-import ml_dtypes
 import numpy as np
 
 if typing.TYPE_CHECKING:
@@ -28,9 +26,7 @@ def pack_4bitx2(array: np.ndarray) -> npt.NDArray[np.uint8]:
     return array_flat[0::2] | array_flat[1::2]  # type: ignore[return-type]
 
 
-def _unpack_uint4_as_uint8(
-    data: npt.NDArray[np.uint8], dims: Sequence[int]
-) -> npt.NDArray[np.uint8]:
+def unpack_4bitx2(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:
@@ -52,56 +48,3 @@ def _unpack_uint4_as_uint8(
         result = result[:-1]
     result.resize(dims, refcheck=False)
     return result
-
-
-def unpack_uint4(
-    data: npt.NDArray[np.uint8], dims: Sequence[int]
-) -> npt.NDArray[ml_dtypes.uint4]:
-    """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.
-    """
-    return _unpack_uint4_as_uint8(data, dims).view(ml_dtypes.uint4)
-
-
-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[ml_dtypes.int4]:
-    """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_as_uint8(data, dims)
-    return _extend_int4_sign_bits(unpacked).view(ml_dtypes.int4)
-
-
-def unpack_float4e2m1(
-    data: npt.NDArray[np.uint8], dims: Sequence[int]
-) -> npt.NDArray[ml_dtypes.float4_e2m1fn]:
-    """Convert a packed float4e2m1 array to unpacked float4e2m1 array.
-
-    Args:
-        data: A numpy array.
-        dims: The dimensions are used to reshape the unpacked buffer.
-
-    Returns:
-        A numpy array of float32 reshaped to dims.
-    """
-    return _unpack_uint4_as_uint8(data, dims).view(ml_dtypes.float4_e2m1fn)

onnx_ir/py.typed

@@ -0,0 +1 @@
+

onnx_ir/serde.py

@@ -74,7 +74,6 @@ from onnx_ir import _convenience, _core, _enums, _protocols, _type_casting
 
 if typing.TYPE_CHECKING:
     import google.protobuf.internal.containers as proto_containers
-    import numpy.typing as npt
 
 logger = logging.getLogger(__name__)
 
@@ -117,13 +116,6 @@ def _little_endian_dtype(dtype) -> np.dtype:
     return np.dtype(dtype).newbyteorder("<")
 
 
-def _unflatten_complex(
-    array: npt.NDArray[np.float32 | np.float64],
-) -> npt.NDArray[np.complex64 | np.complex128]:
-    """Convert the real representation of a complex dtype to the complex dtype."""
-    return array[::2] + 1j * array[1::2]
-
-
 @typing.overload
 def from_proto(proto: onnx.ModelProto) -> _core.Model: ...  # type: ignore[overload-overlap]
 @typing.overload
@@ -391,54 +383,88 @@ class TensorProtoTensor(_core.TensorBase):  # pylint: disable=too-many-ancestors
                 "Cannot convert external tensor to numpy array. Use ir.ExternalTensor instead."
             )
 
+        shape = self._proto.dims
+
         if self._proto.HasField("raw_data"):
-            array = np.frombuffer(self._proto.raw_data, dtype=dtype.numpy().newbyteorder("<"))
-            # Cannot return now, because we may need to unpack 4bit tensors
-        elif dtype == _enums.DataType.STRING:
-            return np.array(self._proto.string_data).reshape(self._proto.dims)
-        elif self._proto.int32_data:
-            array = np.array(self._proto.int32_data, dtype=_little_endian_dtype(np.int32))
-            if dtype in {_enums.DataType.FLOAT16, _enums.DataType.BFLOAT16}:
-                # Reinterpret the int32 as float16 or bfloat16
-                array = array.astype(np.uint16).view(dtype.numpy())
-            elif dtype in {
+            if dtype.bitwidth == 4:
+                return _type_casting.unpack_4bitx2(
+                    np.frombuffer(self._proto.raw_data, dtype=np.uint8), shape
+                ).view(dtype.numpy())
+            return np.frombuffer(
+                self._proto.raw_data, dtype=dtype.numpy().newbyteorder("<")
+            ).reshape(shape)
+        if dtype == _enums.DataType.STRING:
+            return np.array(self._proto.string_data).reshape(shape)
+        if self._proto.int32_data:
+            assert dtype in {
+                _enums.DataType.BFLOAT16,
+                _enums.DataType.BOOL,
+                _enums.DataType.FLOAT16,
+                _enums.DataType.FLOAT4E2M1,
                 _enums.DataType.FLOAT8E4M3FN,
                 _enums.DataType.FLOAT8E4M3FNUZ,
                 _enums.DataType.FLOAT8E5M2,
                 _enums.DataType.FLOAT8E5M2FNUZ,
-            }:
-                array = array.astype(np.uint8).view(dtype.numpy())
-        elif self._proto.int64_data:
-            array = np.array(self._proto.int64_data, dtype=_little_endian_dtype(np.int64))
-        elif self._proto.uint64_data:
+                _enums.DataType.INT16,
+                _enums.DataType.INT32,
+                _enums.DataType.INT4,
+                _enums.DataType.INT8,
+                _enums.DataType.UINT16,
+                _enums.DataType.UINT4,
+                _enums.DataType.UINT8,
+            }, f"Unsupported dtype {dtype} for int32_data"
+            array = np.array(self._proto.int32_data, dtype=_little_endian_dtype(np.int32))
+            if dtype.bitwidth == 32:
+                return array.reshape(shape)
+            if dtype.bitwidth == 16:
+                # Reinterpret the int32 as float16 or bfloat16
+                return array.astype(np.uint16).view(dtype.numpy()).reshape(shape)
+            if dtype.bitwidth == 8:
+                return array.astype(np.uint8).view(dtype.numpy()).reshape(shape)
+            if dtype.bitwidth == 4:
+                return _type_casting.unpack_4bitx2(array.astype(np.uint8), shape).view(
+                    dtype.numpy()
+                )
+            raise ValueError(
+                f"Unsupported dtype {dtype} for int32_data with bitwidth {dtype.bitwidth}"
+            )
+        if self._proto.int64_data:
+            assert dtype in {
+                _enums.DataType.INT64,
+            }, f"Unsupported dtype {dtype} for int64_data"
+            return np.array(
+                self._proto.int64_data, dtype=_little_endian_dtype(np.int64)
+            ).reshape(shape)
+        if self._proto.uint64_data:
+            assert dtype in {
+                _enums.DataType.UINT64,
+                _enums.DataType.UINT32,
+            }, f"Unsupported dtype {dtype} for uint64_data"
             array = np.array(self._proto.uint64_data, dtype=_little_endian_dtype(np.uint64))
-        elif self._proto.float_data:
+            if dtype == _enums.DataType.UINT32:
+                return array.astype(np.uint32).reshape(shape)
+            return array.reshape(shape)
+        if self._proto.float_data:
+            assert dtype in {
+                _enums.DataType.FLOAT,
+                _enums.DataType.COMPLEX64,
+            }, f"Unsupported dtype {dtype} for float_data"
             array = np.array(self._proto.float_data, dtype=_little_endian_dtype(np.float32))
             if dtype == _enums.DataType.COMPLEX64:
-                array = _unflatten_complex(array)
-        elif self._proto.double_data:
+                return array.view(np.complex64).reshape(shape)
+            return array.reshape(shape)
+        if self._proto.double_data:
+            assert dtype in {
+                _enums.DataType.DOUBLE,
+                _enums.DataType.COMPLEX128,
+            }, f"Unsupported dtype {dtype} for double_data"
             array = np.array(self._proto.double_data, dtype=_little_endian_dtype(np.float64))
             if dtype == _enums.DataType.COMPLEX128:
-                array = _unflatten_complex(array)
-        else:
-            # Empty tensor
-            if not self._proto.dims:
-                # When dims not precent and there is no data, we return an empty array
-                return np.array([], dtype=dtype.numpy())
-            else:
-                # Otherwise we return a size 0 array with the correct shape
-                return np.zeros(self._proto.dims, dtype=dtype.numpy())
-
-        if dtype == _enums.DataType.INT4:
-            return _type_casting.unpack_int4(array.astype(np.uint8), self._proto.dims)
-        elif dtype == _enums.DataType.UINT4:
-            return _type_casting.unpack_uint4(array.astype(np.uint8), self._proto.dims)
-        elif dtype == _enums.DataType.FLOAT4E2M1:
-            return _type_casting.unpack_float4e2m1(array.astype(np.uint8), self._proto.dims)
-        else:
-            # Otherwise convert to the correct dtype and reshape
-            # Note we cannot use view() here because the storage dtype may not be the same size as the target
-            return array.astype(dtype.numpy()).reshape(self._proto.dims)
+                return array.view(np.complex128).reshape(shape)
+            return array.reshape(shape)
+
+        # Empty tensor. We return a size 0 array with the correct shape
+        return np.zeros(shape, dtype=dtype.numpy())
 
     def tobytes(self) -> bytes:
         """Return the tensor as a byte string conformed to the ONNX specification, in little endian.

{onnx_ir-0.1.3.dist-info → onnx_ir-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-ir
-Version: 0.1.3
+Version: 0.1.4
 Summary: Efficient in-memory representation for ONNX
 Author-email: ONNX Contributors <onnx-technical-discuss@lists.lfaidata.foundation>
 License: Apache License v2.0

{onnx_ir-0.1.3.dist-info → onnx_ir-0.1.4.dist-info}/RECORD

@@ -1,7 +1,7 @@
-onnx_ir/__init__.py,sha256=5KP1Ngl2qyWiqb5S0Ol5owYsbU0geo4LFwGwN8EXTIk,3424
-onnx_ir/_core.py,sha256=-9BpVTZHuHQ9jsms33wqu4NjMEaDF_M57sIuVxYcM1I,137964
+onnx_ir/__init__.py,sha256=2w65_FPhzimLAzacA-s_IZWCwOWtt6-AMH_Lop7WOcI,3424
+onnx_ir/_core.py,sha256=Y-RJSIgwxKWRTRHc_fWOEH_vjTxtMz2qbZ37hLxzdTI,137284
 onnx_ir/_display.py,sha256=230bMN_hVy47Ug3HkA4o5Tf5Hr21AnBEoq5w0fxjyTs,1300
-onnx_ir/_enums.py,sha256=4lmm_DFKEtz6PqNw6gt6GcqrBYHisctgKMsUbQCm5N8,8252
+onnx_ir/_enums.py,sha256=oWJywcMCMczo1xYz25_eXFOhhJYciDNJTvzgYaKy1E0,9373
 onnx_ir/_graph_comparison.py,sha256=8_D1gu547eCDotEUqxfIJhUGU_Ufhfji7sfsSraOj3g,727
 onnx_ir/_graph_containers.py,sha256=PRKrshRZ5rzWCgRs1TefzJq9n8wyo7OqeKy3XxMhyys,14265
 onnx_ir/_io.py,sha256=GWwA4XOZ-ZX1cgibgaYD0K0O5d9LX21ZwcBN02Wrh04,5205
@@ -11,11 +11,12 @@ onnx_ir/_name_authority.py,sha256=PnoV9TRgMLussZNufWavJXosDWx5avPfldVjMWEEz18,30
 onnx_ir/_polyfill.py,sha256=LzAGBKQbVDlURC0tgQgaxgkYU4rESgCYnqVs-u-Vsx8,887
 onnx_ir/_protocols.py,sha256=M29sIOAvtdlis3QtBvCQPH4pnvSwhJCQNCvs3IrN9FY,21276
 onnx_ir/_tape.py,sha256=nEGY6VZVKuB8FDyXeYr0MTq8j7E4HKOE2yN8qpz4ia0,7007
-onnx_ir/_type_casting.py,sha256=8iZDVrNAx_FwRVt48G4tkzIOFu3I6AsETpH3fdxcyEI,3387
+onnx_ir/_type_casting.py,sha256=hbikTmgFEu0SEfnbgv2R1LbpuPQ2MCfqto3-oLWhcBc,1645
 onnx_ir/_version_utils.py,sha256=bZThuE7meVHFOY1DLsmss9WshVIp9iig7udGfDbVaK4,1333
 onnx_ir/convenience.py,sha256=0B1epuXZCSmY4FbW2vaYfR-t5ubxBZ1UruiytHs-zFw,917
 onnx_ir/external_data.py,sha256=rXHtRU-9tjAt10Iervhr5lsI6Dtv-EhR7J4brxppImA,18079
-onnx_ir/serde.py,sha256=YkbYfQMwn0YAzTd3tVDSWJ-NBiSVsG-74T6xk3e5iTU,75073
+onnx_ir/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+onnx_ir/serde.py,sha256=8D9eBVFcCvwRHyW7Y4CJNuAU0iBI3Mjk0A-w2QidHK4,75892
 onnx_ir/tape.py,sha256=4FyfAHmVhQoMsfHMYnBwP2azi6UF6b6pj--ercObqZs,350
 onnx_ir/tensor_adapters.py,sha256=dXuapwfFcpLhjKC6AOqCXbtY3WvDaEHoCNPwjnUK7_o,6565
 onnx_ir/testing.py,sha256=WTrjf2joWizDWaYMJlV1KjZMQw7YmZ8NvuBTVn1uY6s,8803
@@ -36,8 +37,8 @@ onnx_ir/passes/common/onnx_checker.py,sha256=_sPmJ2ff9pDB1g9q7082BL6fyubomRaj6sv
 onnx_ir/passes/common/shape_inference.py,sha256=LVdvxjeKtcIEbPcb6mKisxoPJOOawzsm3tzk5j9xqeM,3992
 onnx_ir/passes/common/topological_sort.py,sha256=Vcu1YhBdfRX4LROr0NScjB1Pwz2DjBFD0Z_GxqaxPF8,999
 onnx_ir/passes/common/unused_removal.py,sha256=cBNqaqGnUVyCWxsD7hBzYk4qSglVPo3SmHAvkUo5-Oc,7613
-onnx_ir-0.1.3.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-onnx_ir-0.1.3.dist-info/METADATA,sha256=vKG8o_nAUJfjM05rahv0g-FCeHkHXIwCAcuYzSY6PH8,4782
-onnx_ir-0.1.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-onnx_ir-0.1.3.dist-info/top_level.txt,sha256=W5tROO93YjO0XRxIdjMy4wocp-5st5GiI2ukvW7UhDo,8
-onnx_ir-0.1.3.dist-info/RECORD,,
+onnx_ir-0.1.4.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+onnx_ir-0.1.4.dist-info/METADATA,sha256=Oay3Vxf4jfSY50vyCfTYaH0Pbxifv47jd3yimr8CDW8,4782
+onnx_ir-0.1.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+onnx_ir-0.1.4.dist-info/top_level.txt,sha256=W5tROO93YjO0XRxIdjMy4wocp-5st5GiI2ukvW7UhDo,8
+onnx_ir-0.1.4.dist-info/RECORD,,