onnx-ir 0.1.2__tar.gz → 0.1.4__tar.gz

{onnx_ir-0.1.2/src/onnx_ir.egg-info → onnx_ir-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-ir
-Version: 0.1.2
+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.2 → onnx_ir-0.1.4}/src/onnx_ir/__init__.py

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

{onnx_ir-0.1.2 → onnx_ir-0.1.4}/src/onnx_ir/_convenience/__init__.py

@@ -323,6 +323,9 @@ def create_value_mapping(graph: _core.Graph) -> dict[str, _core.Value]:
     and the first value with that name is returned. Values with empty names
     are excluded from the mapping.
 
+    .. versionchanged:: 0.1.2
+        Values from subgraphs are now included in the mapping.
+
     Args:
         graph: The graph to extract the mapping from.
 
@@ -410,6 +413,8 @@ def get_const_tensor(
     it will propagate the shape and type of the constant tensor to the value
     if `propagate_shape_type` is set to True.
 
+    .. versionadded:: 0.1.2
+
     Args:
         value: The value to get the constant tensor from.
         propagate_shape_type: If True, the shape and type of the value will be

{onnx_ir-0.1.2 → onnx_ir-0.1.4}/src/onnx_ir/_core.py

@@ -417,6 +417,9 @@ class Tensor(TensorBase, _protocols.TensorProtocol, Generic[TArrayCompatible]):
         else:
             self._shape = shape
             self._shape.freeze()
+        if isinstance(value, np.generic):
+            # Turn numpy scalar into a numpy array
+            value = np.array(value)  # type: ignore[assignment]
         if dtype is None:
             if isinstance(value, np.ndarray):
                 self._dtype = _enums.DataType.from_numpy(value.dtype)
@@ -654,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,
@@ -672,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)
 
@@ -964,7 +967,10 @@ class LazyTensor(TensorBase, _protocols.TensorProtocol):  # pylint: disable=too-
 
 
 class PackedTensor(TensorBase, _protocols.TensorProtocol, Generic[TArrayCompatible]):  # pylint: disable=too-many-ancestors
-    """A tensor that stores 4bit datatypes in packed format."""
+    """A tensor that stores 4bit datatypes in packed format.
+
+    .. versionadded:: 0.1.2
+    """
 
     __slots__ = (
         "_dtype",
@@ -1065,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."""
@@ -2335,6 +2333,12 @@ class Graph(_protocols.GraphProtocol, Sequence[Node], _display.PrettyPrintable):
     seen as a Sequence of nodes and should be used as such. For example, to obtain
     all nodes as a list, call ``list(graph)``.
 
+    .. versionchanged:: 0.1.1
+        Values with non-none producers will be rejected as graph inputs or initializers.
+
+    .. versionadded:: 0.1.1
+        Added ``add`` method to initializers and attributes.
+
     Attributes:
         name: The name of the graph.
         inputs: The input values of the graph.
@@ -2545,12 +2549,17 @@ class Graph(_protocols.GraphProtocol, Sequence[Node], _display.PrettyPrintable):
         Consider using
         :class:`onnx_ir.traversal.RecursiveGraphIterator` for more advanced
         traversals on nodes.
+
+        .. versionadded:: 0.1.2
         """
         # NOTE: This is a method specific to Graph, not required by the protocol unless proven
         return onnx_ir.traversal.RecursiveGraphIterator(self)
 
     def subgraphs(self) -> Iterator[Graph]:
-        """Get all subgraphs in the graph in O(#nodes + #attributes) time."""
+        """Get all subgraphs in the graph in O(#nodes + #attributes) time.
+
+        .. versionadded:: 0.1.2
+        """
         seen_graphs: set[Graph] = set()
         for node in onnx_ir.traversal.RecursiveGraphIterator(self):
             graph = node.graph
@@ -3216,12 +3225,17 @@ class Function(_protocols.FunctionProtocol, Sequence[Node], _display.PrettyPrint
         Consider using
         :class:`onnx_ir.traversal.RecursiveGraphIterator` for more advanced
         traversals on nodes.
+
+        .. versionadded:: 0.1.2
         """
         # NOTE: This is a method specific to Graph, not required by the protocol unless proven
         return onnx_ir.traversal.RecursiveGraphIterator(self)
 
     def subgraphs(self) -> Iterator[Graph]:
-        """Get all subgraphs in the function in O(#nodes + #attributes) time."""
+        """Get all subgraphs in the function in O(#nodes + #attributes) time.
+
+        .. versionadded:: 0.1.2
+        """
         seen_graphs: set[Graph] = set()
         for node in onnx_ir.traversal.RecursiveGraphIterator(self):
             graph = node.graph

{onnx_ir-0.1.2 → onnx_ir-0.1.4}/src/onnx_ir/_enums.py

@@ -120,6 +120,8 @@ class DataType(enum.IntEnum):
     def bitwidth(self) -> int:
         """Returns the bit width of the data type.
 
+        .. versionadded:: 0.1.2
+
         Raises:
             TypeError: If the data type is not supported.
         """
@@ -167,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-0.1.4/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.2 → onnx_ir-0.1.4}/src/onnx_ir/passes/common/__init__.py

@@ -6,6 +6,7 @@ __all__ = [
     "CheckerPass",
     "ClearMetadataAndDocStringPass",
     "CommonSubexpressionEliminationPass",
+    "DeduplicateInitializersPass",
     "InlinePass",
     "LiftConstantsToInitializersPass",
     "LiftSubgraphInitializersToMainGraphPass",
@@ -29,6 +30,9 @@ from onnx_ir.passes.common.constant_manipulation import (
     LiftSubgraphInitializersToMainGraphPass,
     RemoveInitializersFromInputsPass,
 )
+from onnx_ir.passes.common.initializer_deduplication import (
+    DeduplicateInitializersPass,
+)
 from onnx_ir.passes.common.inliner import InlinePass
 from onnx_ir.passes.common.onnx_checker import CheckerPass
 from onnx_ir.passes.common.shape_inference import ShapeInferencePass

onnx_ir-0.1.4/src/onnx_ir/passes/common/common_subexpression_elimination.py

@@ -0,0 +1,206 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Eliminate common subexpression in ONNX graphs."""
+
+from __future__ import annotations
+
+__all__ = [
+    "CommonSubexpressionEliminationPass",
+]
+
+import logging
+from collections.abc import Sequence
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+class CommonSubexpressionEliminationPass(ir.passes.InPlacePass):
+    """Eliminate common subexpression in ONNX graphs.
+
+    .. versionadded:: 0.1.1
+
+    .. versionchanged:: 0.1.3
+        Constant nodes with values smaller than ``size_limit`` will be CSE'd.
+
+    Attributes:
+        size_limit: The maximum size of the tensor to be csed. If the tensor contains
+            number of elements larger than size_limit, it will not be cse'd. Default is 10.
+
+    """
+
+    def __init__(self, size_limit: int = 10):
+        """Initialize the CommonSubexpressionEliminationPass."""
+        super().__init__()
+        self.size_limit = size_limit
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        """Return the same ir.Model but with CSE applied to the graph."""
+        graph = model.graph
+        modified = self._eliminate_common_subexpression(graph)
+
+        return ir.passes.PassResult(
+            model,
+            modified=modified,
+        )
+
+    def _eliminate_common_subexpression(self, graph: ir.Graph) -> bool:
+        """Eliminate common subexpression in ONNX graphs."""
+        modified: bool = False
+        # node to node identifier, length of outputs, inputs, and attributes
+        existing_node_info_to_the_node: dict[
+            tuple[
+                ir.OperatorIdentifier,
+                int,  # len(outputs)
+                tuple[int, ...],  # input ids
+                tuple[tuple[str, object], ...],  # attributes
+            ],
+            ir.Node,
+        ] = {}
+
+        for node in graph:
+            # Skip control flow ops like Loop and If.
+            control_flow_op: bool = False
+            # Skip large tensors to avoid cse weights and bias.
+            large_tensor: bool = False
+            # Use equality to check if the node is a common subexpression.
+            attributes = {}
+            for k, v in node.attributes.items():
+                # TODO(exporter team): CSE subgraphs.
+                # NOTE: control flow ops like Loop and If won't be CSEd
+                # because attribute: graph won't match.
+                if v.type in (ir.AttributeType.GRAPH, ir.AttributeType.GRAPHS):
+                    control_flow_op = True
+                    break
+                # The attribute value could be directly taken from the original
+                # protobuf, so we need to make a copy of it.
+                value = v.value
+                if v.type in (
+                    ir.AttributeType.INTS,
+                    ir.AttributeType.FLOATS,
+                    ir.AttributeType.STRINGS,
+                ):
+                    # For INT, FLOAT and STRING attributes, we convert them to tuples
+                    # to ensure they are hashable.
+                    value = tuple(value)
+                elif v.type is ir.AttributeType.TENSOR:
+                    if value.size > self.size_limit:
+                        # If the tensor is larger than the size limit, we skip it.
+                        large_tensor = True
+                        break
+                    np_value = value.numpy()
+
+                    value = (np_value.shape, str(np_value.dtype), np_value.tobytes())
+                attributes[k] = value
+
+            if control_flow_op:
+                # If the node is a control flow op, we skip it.
+                logger.debug("Skipping control flow op %s", node)
+                continue
+
+            if large_tensor:
+                # If the node has a large tensor, we skip it.
+                logger.debug("Skipping large tensor in node %s", node)
+                continue
+
+            if _is_non_deterministic_op(node):
+                # If the node is a non-deterministic op, we skip it.
+                logger.debug("Skipping non-deterministic op %s", node)
+                continue
+
+            node_info = (
+                node.op_identifier(),
+                len(node.outputs),
+                tuple(id(input) for input in node.inputs),
+                tuple(sorted(attributes.items())),
+            )
+            # Check if the node is a common subexpression.
+            if node_info in existing_node_info_to_the_node:
+                # If it is, this node has an existing node with the same
+                # operator, number of outputs, inputs, and attributes.
+                # We replace the node with the existing node.
+                modified = True
+                existing_node = existing_node_info_to_the_node[node_info]
+                _remove_node_and_replace_values(
+                    graph,
+                    remove_node=node,
+                    remove_values=node.outputs,
+                    new_values=existing_node.outputs,
+                )
+                logger.debug("Reusing node %s", existing_node)
+            else:
+                # If it is not, add to the mapping.
+                existing_node_info_to_the_node[node_info] = node
+        return modified
+
+
+def _remove_node_and_replace_values(
+    graph: ir.Graph,
+    /,
+    remove_node: ir.Node,
+    remove_values: Sequence[ir.Value],
+    new_values: Sequence[ir.Value],
+) -> None:
+    """Replaces nodes and values in the graph or function.
+
+    Args:
+        graph: The graph to replace nodes and values in.
+        remove_node: The node to remove.
+        remove_values: The values to replace.
+        new_values: The values to replace with.
+    """
+    # Reconnect the users of the deleted values to use the new values
+    ir.convenience.replace_all_uses_with(remove_values, new_values)
+    # Update graph/function outputs if the node generates output
+    if any(remove_value.is_graph_output() for remove_value in remove_values):
+        replacement_mapping = dict(zip(remove_values, new_values))
+        for idx, graph_output in enumerate(graph.outputs):
+            if graph_output in replacement_mapping:
+                new_value = replacement_mapping[graph_output]
+                if new_value.is_graph_output() or new_value.is_graph_input():
+                    # If the new value is also a graph input/output, we need to
+                    # create a Identity node to preserve the remove_value and
+                    # prevent from changing new_value name.
+                    identity_node = ir.node(
+                        "Identity",
+                        inputs=[new_value],
+                        outputs=[
+                            ir.Value(
+                                name=graph_output.name,
+                                type=graph_output.type,
+                                shape=graph_output.shape,
+                            )
+                        ],
+                    )
+                    # reuse the name of the graph output
+                    graph.outputs[idx] = identity_node.outputs[0]
+                    graph.insert_before(
+                        remove_node,
+                        identity_node,
+                    )
+                else:
+                    # if new_value is not graph output, we just
+                    # update it to use old_value name.
+                    new_value.name = graph_output.name
+                    graph.outputs[idx] = new_value
+
+    graph.remove(remove_node, safe=True)
+
+
+def _is_non_deterministic_op(node: ir.Node) -> bool:
+    non_deterministic_ops = frozenset(
+        {
+            "RandomUniform",
+            "RandomNormal",
+            "RandomUniformLike",
+            "RandomNormalLike",
+            "Multinomial",
+        }
+    )
+    return node.op_type in non_deterministic_ops and _is_onnx_domain(node.domain)
+
+
+def _is_onnx_domain(d: str) -> bool:
+    """Check if the domain is the ONNX domain."""
+    return d == ""

onnx_ir-0.1.4/src/onnx_ir/passes/common/initializer_deduplication.py

@@ -0,0 +1,56 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Pass for removing duplicated initializer tensors from a graph."""
+
+from __future__ import annotations
+
+__all__ = [
+    "DeduplicateInitializersPass",
+]
+
+
+import onnx_ir as ir
+
+
+class DeduplicateInitializersPass(ir.passes.InPlacePass):
+    """Remove duplicated initializer tensors from the graph.
+
+    This pass detects initializers with identical shape, dtype, and content,
+    and replaces all duplicate references with a canonical one.
+
+    To deduplicate initializers from subgraphs, use :class:`~onnx_ir.passes.common.LiftSubgraphInitializersToMainGraphPass`
+    to lift the initializers to the main graph first before running pass.
+
+    .. versionadded:: 0.1.3
+    """
+
+    def __init__(self, size_limit: int = 1024):
+        super().__init__()
+        self.size_limit = size_limit
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        graph = model.graph
+        initializers: dict[tuple[ir.DataType, tuple[int, ...], bytes], ir.Value] = {}
+        modified = False
+
+        for initializer in tuple(graph.initializers.values()):
+            # TODO(justinchuby): Handle subgraphs as well. For now users can lift initializers
+            # out from the main graph before running this pass.
+            const_val = initializer.const_value
+            if const_val is None:
+                # Skip if initializer has no constant value
+                continue
+
+            if const_val.size > self.size_limit:
+                continue
+
+            key = (const_val.dtype, tuple(const_val.shape), const_val.tobytes())
+            if key in initializers:
+                modified = True
+                ir.convenience.replace_all_uses_with(initializer, initializers[key])  # type: ignore[index]
+                assert initializer.name is not None
+                graph.initializers.pop(initializer.name)
+            else:
+                initializers[key] = initializer  # type: ignore[index]
+
+        return ir.passes.PassResult(model=model, modified=modified)

onnx_ir-0.1.4/src/onnx_ir/py.typed

@@ -0,0 +1 @@
+

{onnx_ir-0.1.2 → onnx_ir-0.1.4}/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
@@ -200,6 +192,9 @@ def from_onnx_text(
 
     Read more about the textual representation at: https://onnx.ai/onnx/repo-docs/Syntax.html
 
+    .. versionchanged:: 0.1.2
+        Added the ``initializers`` argument.
+
     Args:
         model_text: The ONNX textual representation of the model.
         initializers: Tensors to be added as initializers. If provided, these tensors
@@ -237,6 +232,8 @@ def to_onnx_text(
 ) -> str:
     """Convert the IR model to the ONNX textual representation.
 
+    .. versionadded:: 0.1.2
+
     Args:
         model: The IR model to convert.
         exclude_initializers: If True, the initializers will not be included in the output.
@@ -386,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.2 → onnx_ir-0.1.4}/src/onnx_ir/tensor_adapters.py

@@ -29,6 +29,8 @@ Example::
 from __future__ import annotations
 
 __all__ = [
+    "from_torch_dtype",
+    "to_torch_dtype",
     "TorchTensor",
 ]
 
@@ -44,14 +46,17 @@ if TYPE_CHECKING:
     import torch
 
 
-class TorchTensor(_core.Tensor):
-    def __init__(
-        self, tensor: torch.Tensor, name: str | None = None, doc_string: str | None = None
-    ):
-        # Pass the tensor as the raw data to ir.Tensor's constructor
+_TORCH_DTYPE_TO_ONNX: dict[torch.dtype, ir.DataType] | None = None
+_ONNX_DTYPE_TO_TORCH: dict[ir.DataType, torch.dtype] | None = None
+
+
+def from_torch_dtype(dtype: torch.dtype) -> ir.DataType:
+    """Convert a PyTorch dtype to an ONNX IR DataType."""
+    global _TORCH_DTYPE_TO_ONNX
+    if _TORCH_DTYPE_TO_ONNX is None:
         import torch
 
-        _TORCH_DTYPE_TO_ONNX: dict[torch.dtype, ir.DataType] = {
+        _TORCH_DTYPE_TO_ONNX = {
             torch.bfloat16: ir.DataType.BFLOAT16,
             torch.bool: ir.DataType.BOOL,
             torch.complex128: ir.DataType.COMPLEX128,
@@ -72,8 +77,58 @@ class TorchTensor(_core.Tensor):
             torch.uint32: ir.DataType.UINT32,
             torch.uint64: ir.DataType.UINT64,
         }
+    if dtype not in _TORCH_DTYPE_TO_ONNX:
+        raise TypeError(
+            f"Unsupported PyTorch dtype '{dtype}'. "
+            "Please use a supported dtype from the list: "
+            f"{list(_TORCH_DTYPE_TO_ONNX.keys())}"
+        )
+    return _TORCH_DTYPE_TO_ONNX[dtype]
+
+
+def to_torch_dtype(dtype: ir.DataType) -> torch.dtype:
+    """Convert an ONNX IR DataType to a PyTorch dtype."""
+    global _ONNX_DTYPE_TO_TORCH
+    if _ONNX_DTYPE_TO_TORCH is None:
+        import torch
+
+        _ONNX_DTYPE_TO_TORCH = {
+            ir.DataType.BFLOAT16: torch.bfloat16,
+            ir.DataType.BOOL: torch.bool,
+            ir.DataType.COMPLEX128: torch.complex128,
+            ir.DataType.COMPLEX64: torch.complex64,
+            ir.DataType.FLOAT16: torch.float16,
+            ir.DataType.FLOAT: torch.float32,
+            ir.DataType.DOUBLE: torch.float64,
+            ir.DataType.FLOAT8E4M3FN: torch.float8_e4m3fn,
+            ir.DataType.FLOAT8E4M3FNUZ: torch.float8_e4m3fnuz,
+            ir.DataType.FLOAT8E5M2: torch.float8_e5m2,
+            ir.DataType.FLOAT8E5M2FNUZ: torch.float8_e5m2fnuz,
+            ir.DataType.INT16: torch.int16,
+            ir.DataType.INT32: torch.int32,
+            ir.DataType.INT64: torch.int64,
+            ir.DataType.INT8: torch.int8,
+            ir.DataType.UINT8: torch.uint8,
+            ir.DataType.UINT16: torch.uint16,
+            ir.DataType.UINT32: torch.uint32,
+            ir.DataType.UINT64: torch.uint64,
+        }
+    if dtype not in _ONNX_DTYPE_TO_TORCH:
+        raise TypeError(
+            f"Unsupported conversion from ONNX dtype '{dtype}' to torch. "
+            "Please use a supported dtype from the list: "
+            f"{list(_ONNX_DTYPE_TO_TORCH.keys())}"
+        )
+    return _ONNX_DTYPE_TO_TORCH[dtype]
+
+
+class TorchTensor(_core.Tensor):
+    def __init__(
+        self, tensor: torch.Tensor, name: str | None = None, doc_string: str | None = None
+    ):
+        # Pass the tensor as the raw data to ir.Tensor's constructor
         super().__init__(
-            tensor, dtype=_TORCH_DTYPE_TO_ONNX[tensor.dtype], name=name, doc_string=doc_string
+            tensor, dtype=from_torch_dtype(tensor.dtype), name=name, doc_string=doc_string
         )
 
     def numpy(self) -> npt.NDArray:

{onnx_ir-0.1.2 → onnx_ir-0.1.4/src/onnx_ir.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-ir
-Version: 0.1.2
+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.2 → onnx_ir-0.1.4}/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/initializer_deduplication.py
 src/onnx_ir/passes/common/inliner.py
 src/onnx_ir/passes/common/onnx_checker.py
 src/onnx_ir/passes/common/shape_inference.py

onnx_ir-0.1.2/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)

onnx_ir-0.1.2/src/onnx_ir/passes/common/common_subexpression_elimination.py

@@ -1,177 +0,0 @@
-# Copyright (c) ONNX Project Contributors
-# SPDX-License-Identifier: Apache-2.0
-"""Eliminate common subexpression in ONNX graphs."""
-
-from __future__ import annotations
-
-__all__ = [
-    "CommonSubexpressionEliminationPass",
-]
-
-import logging
-from collections.abc import Sequence
-
-import onnx_ir as ir
-
-logger = logging.getLogger(__name__)
-
-
-class CommonSubexpressionEliminationPass(ir.passes.InPlacePass):
-    """Eliminate common subexpression in ONNX graphs."""
-
-    def call(self, model: ir.Model) -> ir.passes.PassResult:
-        """Return the same ir.Model but with CSE applied to the graph."""
-        modified = False
-        graph = model.graph
-
-        modified = _eliminate_common_subexpression(graph, modified)
-
-        return ir.passes.PassResult(
-            model,
-            modified=modified,
-        )
-
-
-def _eliminate_common_subexpression(graph: ir.Graph, modified: bool) -> bool:
-    """Eliminate common subexpression in ONNX graphs."""
-    # node to node identifier, length of outputs, inputs, and attributes
-    existing_node_info_to_the_node: dict[
-        tuple[
-            ir.OperatorIdentifier,
-            int,  # len(outputs)
-            tuple[int, ...],  # input ids
-            tuple[tuple[str, object], ...],  # attributes
-        ],
-        ir.Node,
-    ] = {}
-
-    for node in graph:
-        # Skip control flow ops like Loop and If.
-        control_flow_op: bool = False
-        # Use equality to check if the node is a common subexpression.
-        attributes = {}
-        for k, v in node.attributes.items():
-            # TODO(exporter team): CSE subgraphs.
-            # NOTE: control flow ops like Loop and If won't be CSEd
-            # because attribute: graph won't match.
-            if v.type in (ir.AttributeType.GRAPH, ir.AttributeType.GRAPHS):
-                control_flow_op = True
-                logger.debug("Skipping control flow op %s", node)
-            # The attribute value could be directly taken from the original
-            # protobuf, so we need to make a copy of it.
-            value = v.value
-            if v.type in (
-                ir.AttributeType.INTS,
-                ir.AttributeType.FLOATS,
-                ir.AttributeType.STRINGS,
-            ):
-                # For INT, FLOAT and STRING attributes, we convert them to tuples
-                # to ensure they are hashable.
-                value = tuple(value)
-            attributes[k] = value
-
-        if control_flow_op:
-            # If the node is a control flow op, we skip it.
-            logger.debug("Skipping control flow op %s", node)
-            continue
-
-        if _is_non_deterministic_op(node):
-            # If the node is a non-deterministic op, we skip it.
-            logger.debug("Skipping non-deterministic op %s", node)
-            continue
-
-        node_info = (
-            node.op_identifier(),
-            len(node.outputs),
-            tuple(id(input) for input in node.inputs),
-            tuple(sorted(attributes.items())),
-        )
-        # Check if the node is a common subexpression.
-        if node_info in existing_node_info_to_the_node:
-            # If it is, this node has an existing node with the same
-            # operator, number of outputs, inputs, and attributes.
-            # We replace the node with the existing node.
-            modified = True
-            existing_node = existing_node_info_to_the_node[node_info]
-            _remove_node_and_replace_values(
-                graph,
-                remove_node=node,
-                remove_values=node.outputs,
-                new_values=existing_node.outputs,
-            )
-            logger.debug("Reusing node %s", existing_node)
-        else:
-            # If it is not, add to the mapping.
-            existing_node_info_to_the_node[node_info] = node
-    return modified
-
-
-def _remove_node_and_replace_values(
-    graph: ir.Graph,
-    /,
-    remove_node: ir.Node,
-    remove_values: Sequence[ir.Value],
-    new_values: Sequence[ir.Value],
-) -> None:
-    """Replaces nodes and values in the graph or function.
-
-    Args:
-        graph: The graph to replace nodes and values in.
-        remove_node: The node to remove.
-        remove_values: The values to replace.
-        new_values: The values to replace with.
-    """
-    # Reconnect the users of the deleted values to use the new values
-    ir.convenience.replace_all_uses_with(remove_values, new_values)
-    # Update graph/function outputs if the node generates output
-    if any(remove_value.is_graph_output() for remove_value in remove_values):
-        replacement_mapping = dict(zip(remove_values, new_values))
-        for idx, graph_output in enumerate(graph.outputs):
-            if graph_output in replacement_mapping:
-                new_value = replacement_mapping[graph_output]
-                if new_value.is_graph_output() or new_value.is_graph_input():
-                    # If the new value is also a graph input/output, we need to
-                    # create a Identity node to preserve the remove_value and
-                    # prevent from changing new_value name.
-                    identity_node = ir.node(
-                        "Identity",
-                        inputs=[new_value],
-                        outputs=[
-                            ir.Value(
-                                name=graph_output.name,
-                                type=graph_output.type,
-                                shape=graph_output.shape,
-                            )
-                        ],
-                    )
-                    # reuse the name of the graph output
-                    graph.outputs[idx] = identity_node.outputs[0]
-                    graph.insert_before(
-                        remove_node,
-                        identity_node,
-                    )
-                else:
-                    # if new_value is not graph output, we just
-                    # update it to use old_value name.
-                    new_value.name = graph_output.name
-                    graph.outputs[idx] = new_value
-
-    graph.remove(remove_node, safe=True)
-
-
-def _is_non_deterministic_op(node: ir.Node) -> bool:
-    non_deterministic_ops = frozenset(
-        {
-            "RandomUniform",
-            "RandomNormal",
-            "RandomUniformLike",
-            "RandomNormalLike",
-            "Multinomial",
-        }
-    )
-    return node.op_type in non_deterministic_ops and _is_onnx_domain(node.domain)
-
-
-def _is_onnx_domain(d: str) -> bool:
-    """Check if the domain is the ONNX domain."""
-    return d == ""