onnx-ir 0.1.3__tar.gz → 0.1.5__tar.gz

{onnx_ir-0.1.3/src/onnx_ir.egg-info → onnx_ir-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-ir
-Version: 0.1.3
+Version: 0.1.5
 Summary: Efficient in-memory representation for ONNX
 Author-email: ONNX Contributors <onnx-technical-discuss@lists.lfaidata.foundation>
 License: Apache License v2.0
@@ -30,6 +30,7 @@ Dynamic: license-file
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![codecov](https://codecov.io/gh/onnx/ir-py/graph/badge.svg?token=SPQ3G9T78Z)](https://codecov.io/gh/onnx/ir-py)
 [![DeepWiki](https://img.shields.io/badge/DeepWiki-onnx%2Fir--py-blue.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACwAAAAyCAYAAAAnWDnqAAAAAXNSR0IArs4c6QAAA05JREFUaEPtmUtyEzEQhtWTQyQLHNak2AB7ZnyXZMEjXMGeK/AIi+QuHrMnbChYY7MIh8g01fJoopFb0uhhEqqcbWTp06/uv1saEDv4O3n3dV60RfP947Mm9/SQc0ICFQgzfc4CYZoTPAswgSJCCUJUnAAoRHOAUOcATwbmVLWdGoH//PB8mnKqScAhsD0kYP3j/Yt5LPQe2KvcXmGvRHcDnpxfL2zOYJ1mFwrryWTz0advv1Ut4CJgf5uhDuDj5eUcAUoahrdY/56ebRWeraTjMt/00Sh3UDtjgHtQNHwcRGOC98BJEAEymycmYcWwOprTgcB6VZ5JK5TAJ+fXGLBm3FDAmn6oPPjR4rKCAoJCal2eAiQp2x0vxTPB3ALO2CRkwmDy5WohzBDwSEFKRwPbknEggCPB/imwrycgxX2NzoMCHhPkDwqYMr9tRcP5qNrMZHkVnOjRMWwLCcr8ohBVb1OMjxLwGCvjTikrsBOiA6fNyCrm8V1rP93iVPpwaE+gO0SsWmPiXB+jikdf6SizrT5qKasx5j8ABbHpFTx+vFXp9EnYQmLx02h1QTTrl6eDqxLnGjporxl3NL3agEvXdT0WmEost648sQOYAeJS9Q7bfUVoMGnjo4AZdUMQku50McDcMWcBPvr0SzbTAFDfvJqwLzgxwATnCgnp4wDl6Aa+Ax283gghmj+vj7feE2KBBRMW3FzOpLOADl0Isb5587h/U4gGvkt5v60Z1VLG8BhYjbzRwyQZemwAd6cCR5/XFWLYZRIMpX39AR0tjaGGiGzLVyhse5C9RKC6ai42ppWPKiBagOvaYk8lO7DajerabOZP46Lby5wKjw1HCRx7p9sVMOWGzb/vA1hwiWc6jm3MvQDTogQkiqIhJV0nBQBTU+3okKCFDy9WwferkHjtxib7t3xIUQtHxnIwtx4mpg26/HfwVNVDb4oI9RHmx5WGelRVlrtiw43zboCLaxv46AZeB3IlTkwouebTr1y2NjSpHz68WNFjHvupy3q8TFn3Hos2IAk4Ju5dCo8B3wP7VPr/FGaKiG+T+v+TQqIrOqMTL1VdWV1DdmcbO8KXBz6esmYWYKPwDL5b5FA1a0hwapHiom0r/cKaoqr+27/XcrS5UwSMbQAAAABJRU5ErkJggg==)](https://deepwiki.com/onnx/ir-py)
+[![PyPI Downloads](https://static.pepy.tech/badge/onnx-ir/month)](https://pepy.tech/projects/onnx-ir)
 
 An in-memory IR that supports the full ONNX spec, designed for graph construction, analysis and transformation.
 

{onnx_ir-0.1.3 → onnx_ir-0.1.5}/README.md

@@ -5,6 +5,7 @@
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![codecov](https://codecov.io/gh/onnx/ir-py/graph/badge.svg?token=SPQ3G9T78Z)](https://codecov.io/gh/onnx/ir-py)
 [![DeepWiki](https://img.shields.io/badge/DeepWiki-onnx%2Fir--py-blue.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACwAAAAyCAYAAAAnWDnqAAAAAXNSR0IArs4c6QAAA05JREFUaEPtmUtyEzEQhtWTQyQLHNak2AB7ZnyXZMEjXMGeK/AIi+QuHrMnbChYY7MIh8g01fJoopFb0uhhEqqcbWTp06/uv1saEDv4O3n3dV60RfP947Mm9/SQc0ICFQgzfc4CYZoTPAswgSJCCUJUnAAoRHOAUOcATwbmVLWdGoH//PB8mnKqScAhsD0kYP3j/Yt5LPQe2KvcXmGvRHcDnpxfL2zOYJ1mFwrryWTz0advv1Ut4CJgf5uhDuDj5eUcAUoahrdY/56ebRWeraTjMt/00Sh3UDtjgHtQNHwcRGOC98BJEAEymycmYcWwOprTgcB6VZ5JK5TAJ+fXGLBm3FDAmn6oPPjR4rKCAoJCal2eAiQp2x0vxTPB3ALO2CRkwmDy5WohzBDwSEFKRwPbknEggCPB/imwrycgxX2NzoMCHhPkDwqYMr9tRcP5qNrMZHkVnOjRMWwLCcr8ohBVb1OMjxLwGCvjTikrsBOiA6fNyCrm8V1rP93iVPpwaE+gO0SsWmPiXB+jikdf6SizrT5qKasx5j8ABbHpFTx+vFXp9EnYQmLx02h1QTTrl6eDqxLnGjporxl3NL3agEvXdT0WmEost648sQOYAeJS9Q7bfUVoMGnjo4AZdUMQku50McDcMWcBPvr0SzbTAFDfvJqwLzgxwATnCgnp4wDl6Aa+Ax283gghmj+vj7feE2KBBRMW3FzOpLOADl0Isb5587h/U4gGvkt5v60Z1VLG8BhYjbzRwyQZemwAd6cCR5/XFWLYZRIMpX39AR0tjaGGiGzLVyhse5C9RKC6ai42ppWPKiBagOvaYk8lO7DajerabOZP46Lby5wKjw1HCRx7p9sVMOWGzb/vA1hwiWc6jm3MvQDTogQkiqIhJV0nBQBTU+3okKCFDy9WwferkHjtxib7t3xIUQtHxnIwtx4mpg26/HfwVNVDb4oI9RHmx5WGelRVlrtiw43zboCLaxv46AZeB3IlTkwouebTr1y2NjSpHz68WNFjHvupy3q8TFn3Hos2IAk4Ju5dCo8B3wP7VPr/FGaKiG+T+v+TQqIrOqMTL1VdWV1DdmcbO8KXBz6esmYWYKPwDL5b5FA1a0hwapHiom0r/cKaoqr+27/XcrS5UwSMbQAAAABJRU5ErkJggg==)](https://deepwiki.com/onnx/ir-py)
+[![PyPI Downloads](https://static.pepy.tech/badge/onnx-ir/month)](https://pepy.tech/projects/onnx-ir)
 
 An in-memory IR that supports the full ONNX spec, designed for graph construction, analysis and transformation.
 

{onnx_ir-0.1.3 → onnx_ir-0.1.5}/src/onnx_ir/__init__.py

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

{onnx_ir-0.1.3 → onnx_ir-0.1.5}/src/onnx_ir/_core.py

@@ -78,6 +78,7 @@ _NON_NUMPY_NATIVE_TYPES = frozenset(
         _enums.DataType.FLOAT8E4M3FNUZ,
         _enums.DataType.FLOAT8E5M2,
         _enums.DataType.FLOAT8E5M2FNUZ,
+        _enums.DataType.FLOAT8E8M0,
         _enums.DataType.INT4,
         _enums.DataType.UINT4,
         _enums.DataType.FLOAT4E2M1,
@@ -261,6 +262,7 @@ def _check_numpy_representation_type(array: np.ndarray, dtype: _enums.DataType)
             ml_dtypes.float8_e4m3fn,
             ml_dtypes.float8_e5m2fnuz,
             ml_dtypes.float8_e5m2,
+            ml_dtypes.float8_e8m0fnu,
         ):
             raise TypeError(
                 f"The numpy array dtype must be uint8 or ml_dtypes.float8* (not {array.dtype}) for IR data type {dtype}."
@@ -319,6 +321,8 @@ def _maybe_view_np_array_with_ml_dtypes(
         return array.view(ml_dtypes.float8_e5m2)
     if dtype == _enums.DataType.FLOAT8E5M2FNUZ:
         return array.view(ml_dtypes.float8_e5m2fnuz)
+    if dtype == _enums.DataType.FLOAT8E8M0:
+        return array.view(ml_dtypes.float8_e8m0fnu)
     if dtype == _enums.DataType.INT4:
         return array.view(ml_dtypes.int4)
     if dtype == _enums.DataType.UINT4:
@@ -657,15 +661,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 +677,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 +1075,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-0.1.3 → onnx_ir-0.1.5}/src/onnx_ir/_enums.py

@@ -65,6 +65,7 @@ class DataType(enum.IntEnum):
     UINT4 = 21
     INT4 = 22
     FLOAT4E2M1 = 23
+    FLOAT8E8M0 = 24
 
     @classmethod
     def from_numpy(cls, dtype: np.dtype) -> DataType:
@@ -81,6 +82,7 @@ class DataType(enum.IntEnum):
 
         # Special cases for handling custom dtypes defined in ONNX (as of onnx 1.18)
         # Ref: https://github.com/onnx/onnx/blob/2d42b6a60a52e925e57c422593e88cc51890f58a/onnx/_custom_element_types.py
+        # TODO(#137): Remove this when ONNX 1.19 is the minimum requirement
         if hasattr(dtype, "names"):
             if dtype.names == ("bfloat16",):
                 return DataType.BFLOAT16
@@ -167,6 +169,50 @@ class DataType(enum.IntEnum):
             DataType.FLOAT8E5M2,
             DataType.FLOAT8E5M2FNUZ,
             DataType.FLOAT4E2M1,
+            DataType.FLOAT8E8M0,
+        }
+
+    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,
+            DataType.FLOAT8E8M0,
         }
 
     def __repr__(self) -> str:
@@ -199,6 +245,7 @@ _BITWIDTH_MAP = {
     DataType.UINT4: 4,
     DataType.INT4: 4,
     DataType.FLOAT4E2M1: 4,
+    DataType.FLOAT8E8M0: 8,
 }
 
 
@@ -224,6 +271,7 @@ _NP_TYPE_TO_DATA_TYPE = {
     np.dtype(ml_dtypes.float8_e4m3fnuz): DataType.FLOAT8E4M3FNUZ,
     np.dtype(ml_dtypes.float8_e5m2): DataType.FLOAT8E5M2,
     np.dtype(ml_dtypes.float8_e5m2fnuz): DataType.FLOAT8E5M2FNUZ,
+    np.dtype(ml_dtypes.float8_e8m0fnu): DataType.FLOAT8E8M0,
     np.dtype(ml_dtypes.int4): DataType.INT4,
     np.dtype(ml_dtypes.uint4): DataType.UINT4,
 }
@@ -248,6 +296,7 @@ _DATA_TYPE_TO_SHORT_NAME = {
     DataType.FLOAT8E5M2: "f8e5m2",
     DataType.FLOAT8E4M3FNUZ: "f8e4m3fnuz",
     DataType.FLOAT8E5M2FNUZ: "f8e5m2fnuz",
+    DataType.FLOAT8E8M0: "f8e8m0",
     DataType.FLOAT4E2M1: "f4e2m1",
     DataType.COMPLEX64: "c64",
     DataType.COMPLEX128: "c128",

onnx_ir-0.1.5/src/onnx_ir/_type_casting.py

@@ -0,0 +1,50 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Numpy utilities for non-native type operation."""
+
+from __future__ import annotations
+
+import typing
+from collections.abc import Sequence
+
+import numpy as np
+
+if typing.TYPE_CHECKING:
+    import numpy.typing as npt
+
+
+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_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:
+        data: A numpy array.
+        dims: The dimensions are used to reshape the unpacked buffer.
+
+    Returns:
+        A numpy array of int8/uint8 reshaped to dims.
+    """
+    assert data.dtype == np.uint8, "Input data must be of type uint8"
+    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

{onnx_ir-0.1.3 → onnx_ir-0.1.5}/src/onnx_ir/passes/common/__init__.py

@@ -7,6 +7,7 @@ __all__ = [
     "ClearMetadataAndDocStringPass",
     "CommonSubexpressionEliminationPass",
     "DeduplicateInitializersPass",
+    "IdentityEliminationPass",
     "InlinePass",
     "LiftConstantsToInitializersPass",
     "LiftSubgraphInitializersToMainGraphPass",
@@ -30,6 +31,9 @@ from onnx_ir.passes.common.constant_manipulation import (
     LiftSubgraphInitializersToMainGraphPass,
     RemoveInitializersFromInputsPass,
 )
+from onnx_ir.passes.common.identity_elimination import (
+    IdentityEliminationPass,
+)
 from onnx_ir.passes.common.initializer_deduplication import (
     DeduplicateInitializersPass,
 )

onnx_ir-0.1.5/src/onnx_ir/passes/common/identity_elimination.py

@@ -0,0 +1,97 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Identity elimination pass for removing redundant Identity nodes."""
+
+from __future__ import annotations
+
+__all__ = [
+    "IdentityEliminationPass",
+]
+
+import logging
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+class IdentityEliminationPass(ir.passes.InPlacePass):
+    """Pass for eliminating redundant Identity nodes.
+
+    This pass removes Identity nodes according to the following rules:
+    1. For any node of the form `y = Identity(x)`, where `y` is not an output
+       of any graph, replace all uses of `y` with a use of `x`, and remove the node.
+    2. If `y` is an output of a graph, and `x` is not an input of any graph,
+       we can still do the elimination, but the value `x` should be renamed to be `y`.
+    3. If `y` is a graph-output and `x` is a graph-input, we cannot eliminate
+       the node. It should be retained.
+    """
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        """Main entry point for the identity elimination pass."""
+        modified = False
+
+        # Use RecursiveGraphIterator to process all nodes in the model graph and subgraphs
+        for node in ir.traversal.RecursiveGraphIterator(model.graph):
+            if self._try_eliminate_identity_node(node):
+                modified = True
+
+        # Process nodes in functions
+        for function in model.functions.values():
+            for node in ir.traversal.RecursiveGraphIterator(function):
+                if self._try_eliminate_identity_node(node):
+                    modified = True
+
+        if modified:
+            logger.info("Identity elimination pass modified the model")
+
+        return ir.passes.PassResult(model, modified=modified)
+
+    def _try_eliminate_identity_node(self, node: ir.Node) -> bool:
+        """Try to eliminate a single identity node. Returns True if modified."""
+        if node.op_type != "Identity" or node.domain != "":
+            return False
+
+        if len(node.inputs) != 1 or len(node.outputs) != 1:
+            # Invalid Identity node, skip
+            return False
+
+        input_value = node.inputs[0]
+        output_value = node.outputs[0]
+
+        if input_value is None:
+            # Cannot eliminate if input is None
+            return False
+
+        # Get the graph that contains this node
+        graph_like = node.graph
+        assert graph_like is not None, "Node must be in a graph"
+
+        output_is_graph_output = output_value.is_graph_output()
+        input_is_graph_input = input_value.is_graph_input()
+
+        # Case 3: Both output is graph output and input is graph input - keep the node
+        if output_is_graph_output and input_is_graph_input:
+            return False
+
+        # Case 1 & 2 (merged): Eliminate the identity node
+        # Replace all uses of output with input
+        ir.convenience.replace_all_uses_with(output_value, input_value)
+
+        # If output is a graph output, we need to rename input and update graph outputs
+        if output_is_graph_output:
+            # Store the original output name
+            original_output_name = output_value.name
+
+            # Update the input value to have the output's name
+            input_value.name = original_output_name
+
+            # Update graph outputs to point to the input value
+            for idx, graph_output in enumerate(graph_like.outputs):
+                if graph_output is output_value:
+                    graph_like.outputs[idx] = input_value
+
+        # Remove the identity node
+        graph_like.remove(node, safe=True)
+        logger.debug("Eliminated identity node: %s", node)
+        return True

onnx_ir-0.1.5/src/onnx_ir/py.typed

@@ -0,0 +1 @@
+

{onnx_ir-0.1.3 → onnx_ir-0.1.5}/src/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,89 @@ 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.FLOAT8E8M0,
+                _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.
@@ -479,6 +506,7 @@ class TensorProtoTensor(_core.TensorBase):  # pylint: disable=too-many-ancestors
                 _enums.DataType.FLOAT8E4M3FNUZ,
                 _enums.DataType.FLOAT8E5M2,
                 _enums.DataType.FLOAT8E5M2FNUZ,
+                _enums.DataType.FLOAT8E8M0,
                 _enums.DataType.INT4,
                 _enums.DataType.UINT4,
                 _enums.DataType.FLOAT4E2M1,

{onnx_ir-0.1.3 → onnx_ir-0.1.5}/src/onnx_ir/tensor_adapters.py

@@ -68,6 +68,7 @@ def from_torch_dtype(dtype: torch.dtype) -> ir.DataType:
             torch.float8_e4m3fnuz: ir.DataType.FLOAT8E4M3FNUZ,
             torch.float8_e5m2: ir.DataType.FLOAT8E5M2,
             torch.float8_e5m2fnuz: ir.DataType.FLOAT8E5M2FNUZ,
+            torch.float8_e8m0fnu: ir.DataType.FLOAT8E8M0,
             torch.int16: ir.DataType.INT16,
             torch.int32: ir.DataType.INT32,
             torch.int64: ir.DataType.INT64,
@@ -104,6 +105,7 @@ def to_torch_dtype(dtype: ir.DataType) -> torch.dtype:
             ir.DataType.FLOAT8E4M3FNUZ: torch.float8_e4m3fnuz,
             ir.DataType.FLOAT8E5M2: torch.float8_e5m2,
             ir.DataType.FLOAT8E5M2FNUZ: torch.float8_e5m2fnuz,
+            ir.DataType.FLOAT8E8M0: torch.float8_e8m0fnu,
             ir.DataType.INT16: torch.int16,
             ir.DataType.INT32: torch.int32,
             ir.DataType.INT64: torch.int64,
@@ -142,6 +144,7 @@ class TorchTensor(_core.Tensor):
             ir.DataType.FLOAT8E4M3FNUZ,
             ir.DataType.FLOAT8E5M2,
             ir.DataType.FLOAT8E5M2FNUZ,
+            ir.DataType.FLOAT8E8M0,
         }:
             return self.raw.view(torch.uint8).numpy(force=True).view(self.dtype.numpy())
 

{onnx_ir-0.1.3 → onnx_ir-0.1.5/src/onnx_ir.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-ir
-Version: 0.1.3
+Version: 0.1.5
 Summary: Efficient in-memory representation for ONNX
 Author-email: ONNX Contributors <onnx-technical-discuss@lists.lfaidata.foundation>
 License: Apache License v2.0
@@ -30,6 +30,7 @@ Dynamic: license-file
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![codecov](https://codecov.io/gh/onnx/ir-py/graph/badge.svg?token=SPQ3G9T78Z)](https://codecov.io/gh/onnx/ir-py)
 [![DeepWiki](https://img.shields.io/badge/DeepWiki-onnx%2Fir--py-blue.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACwAAAAyCAYAAAAnWDnqAAAAAXNSR0IArs4c6QAAA05JREFUaEPtmUtyEzEQhtWTQyQLHNak2AB7ZnyXZMEjXMGeK/AIi+QuHrMnbChYY7MIh8g01fJoopFb0uhhEqqcbWTp06/uv1saEDv4O3n3dV60RfP947Mm9/SQc0ICFQgzfc4CYZoTPAswgSJCCUJUnAAoRHOAUOcATwbmVLWdGoH//PB8mnKqScAhsD0kYP3j/Yt5LPQe2KvcXmGvRHcDnpxfL2zOYJ1mFwrryWTz0advv1Ut4CJgf5uhDuDj5eUcAUoahrdY/56ebRWeraTjMt/00Sh3UDtjgHtQNHwcRGOC98BJEAEymycmYcWwOprTgcB6VZ5JK5TAJ+fXGLBm3FDAmn6oPPjR4rKCAoJCal2eAiQp2x0vxTPB3ALO2CRkwmDy5WohzBDwSEFKRwPbknEggCPB/imwrycgxX2NzoMCHhPkDwqYMr9tRcP5qNrMZHkVnOjRMWwLCcr8ohBVb1OMjxLwGCvjTikrsBOiA6fNyCrm8V1rP93iVPpwaE+gO0SsWmPiXB+jikdf6SizrT5qKasx5j8ABbHpFTx+vFXp9EnYQmLx02h1QTTrl6eDqxLnGjporxl3NL3agEvXdT0WmEost648sQOYAeJS9Q7bfUVoMGnjo4AZdUMQku50McDcMWcBPvr0SzbTAFDfvJqwLzgxwATnCgnp4wDl6Aa+Ax283gghmj+vj7feE2KBBRMW3FzOpLOADl0Isb5587h/U4gGvkt5v60Z1VLG8BhYjbzRwyQZemwAd6cCR5/XFWLYZRIMpX39AR0tjaGGiGzLVyhse5C9RKC6ai42ppWPKiBagOvaYk8lO7DajerabOZP46Lby5wKjw1HCRx7p9sVMOWGzb/vA1hwiWc6jm3MvQDTogQkiqIhJV0nBQBTU+3okKCFDy9WwferkHjtxib7t3xIUQtHxnIwtx4mpg26/HfwVNVDb4oI9RHmx5WGelRVlrtiw43zboCLaxv46AZeB3IlTkwouebTr1y2NjSpHz68WNFjHvupy3q8TFn3Hos2IAk4Ju5dCo8B3wP7VPr/FGaKiG+T+v+TQqIrOqMTL1VdWV1DdmcbO8KXBz6esmYWYKPwDL5b5FA1a0hwapHiom0r/cKaoqr+27/XcrS5UwSMbQAAAABJRU5ErkJggg==)](https://deepwiki.com/onnx/ir-py)
+[![PyPI Downloads](https://static.pepy.tech/badge/onnx-ir/month)](https://pepy.tech/projects/onnx-ir)
 
 An in-memory IR that supports the full ONNX spec, designed for graph construction, analysis and transformation.
 

{onnx_ir-0.1.3 → onnx_ir-0.1.5}/src/onnx_ir.egg-info/SOURCES.txt

@@ -19,6 +19,7 @@ src/onnx_ir/_type_casting.py
 src/onnx_ir/_version_utils.py
 src/onnx_ir/convenience.py
 src/onnx_ir/external_data.py
+src/onnx_ir/py.typed
 src/onnx_ir/serde.py
 src/onnx_ir/tape.py
 src/onnx_ir/tensor_adapters.py
@@ -39,6 +40,7 @@ src/onnx_ir/passes/common/_c_api_utils.py
 src/onnx_ir/passes/common/clear_metadata_and_docstring.py
 src/onnx_ir/passes/common/common_subexpression_elimination.py
 src/onnx_ir/passes/common/constant_manipulation.py
+src/onnx_ir/passes/common/identity_elimination.py
 src/onnx_ir/passes/common/initializer_deduplication.py
 src/onnx_ir/passes/common/inliner.py
 src/onnx_ir/passes/common/onnx_checker.py

onnx_ir-0.1.3/src/onnx_ir/_type_casting.py

@@ -1,107 +0,0 @@
-# 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:
-    import numpy.typing as npt
-
-
-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_uint4_as_uint8(
-    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.
-    """
-    assert data.dtype == np.uint8, "Input data must be of type uint8"
-    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 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)