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

onnx_ir/_type_casting.py

@@ -15,7 +15,7 @@ if typing.TYPE_CHECKING:
     import numpy.typing as npt
 
 
-def pack_int4(array: np.ndarray) -> npt.NDArray[np.uint8]:
+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()
@@ -40,6 +40,7 @@ def _unpack_uint4_as_uint8(
     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)

onnx_ir/_version_utils.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """Version utils for testing."""
 
+# pylint: disable=import-outside-toplevel
 from __future__ import annotations
 
 import packaging.version
@@ -9,7 +10,7 @@ import packaging.version
 
 def onnx_older_than(version: str) -> bool:
     """Returns True if the ONNX version is older than the given version."""
-    import onnx  # pylint: disable=import-outside-toplevel
+    import onnx  # noqa: TID251
 
     return (
         packaging.version.parse(onnx.__version__).release
@@ -19,7 +20,7 @@ def onnx_older_than(version: str) -> bool:
 
 def torch_older_than(version: str) -> bool:
     """Returns True if the torch version is older than the given version."""
-    import torch  # pylint: disable=import-outside-toplevel
+    import torch
 
     return (
         packaging.version.parse(torch.__version__).release
@@ -27,42 +28,9 @@ def torch_older_than(version: str) -> bool:
     )
 
 
-def transformers_older_than(version: str) -> bool | None:
-    """Returns True if the transformers version is older than the given version."""
-    try:
-        import transformers  # pylint: disable=import-outside-toplevel
-    except ImportError:
-        return None
-
-    return (
-        packaging.version.parse(transformers.__version__).release
-        < packaging.version.parse(version).release
-    )
-
-
-def is_onnxruntime_training() -> bool:
-    """Returns True if the onnxruntime is onnxruntime-training."""
-    try:
-        from onnxruntime import training  # pylint: disable=import-outside-toplevel
-
-        assert training
-    except ImportError:
-        # onnxruntime not training
-        return False
-
-    try:
-        from onnxruntime.capi.onnxruntime_pybind11_state import (  # pylint: disable=import-outside-toplevel
-            OrtValueVector,
-        )
-    except ImportError:
-        return False
-
-    return hasattr(OrtValueVector, "push_back_batch")
-
-
 def onnxruntime_older_than(version: str) -> bool:
     """Returns True if the onnxruntime version is older than the given version."""
-    import onnxruntime  # pylint: disable=import-outside-toplevel
+    import onnxruntime
 
     return (
         packaging.version.parse(onnxruntime.__version__).release
@@ -72,20 +40,9 @@ def onnxruntime_older_than(version: str) -> bool:
 
 def numpy_older_than(version: str) -> bool:
     """Returns True if the numpy version is older than the given version."""
-    import numpy  # pylint: disable=import-outside-toplevel
+    import numpy
 
     return (
         packaging.version.parse(numpy.__version__).release
         < packaging.version.parse(version).release
     )
-
-
-def has_transformers():
-    """Tells if transformers is installed."""
-    try:
-        import transformers  # pylint: disable=import-outside-toplevel
-
-        assert transformers
-        return True  # noqa
-    except ImportError:
-        return False

onnx_ir/convenience.py

@@ -7,15 +7,17 @@ from __future__ import annotations
 __all__ = [
     "convert_attribute",
     "convert_attributes",
+    "create_value_mapping",
+    "get_const_tensor",
     "replace_all_uses_with",
     "replace_nodes_and_values",
-    "create_value_mapping",
 ]
 
 from onnx_ir._convenience import (
     convert_attribute,
     convert_attributes,
     create_value_mapping,
+    get_const_tensor,
     replace_all_uses_with,
     replace_nodes_and_values,
 )

onnx_ir/external_data.py

@@ -4,12 +4,15 @@
 
 from __future__ import annotations
 
+from typing import Callable
+
 __all__ = [
     "set_base_dir",
     "unload_from_model",
     "load_to_model",
     "convert_tensors_to_external",
     "convert_tensors_from_external",
+    "CallbackInfo",
 ]
 
 import dataclasses
@@ -48,6 +51,21 @@ class _ExternalDataInfo:
     length: int
 
 
+@dataclasses.dataclass
+class CallbackInfo:
+    """A class that shares information about a tensor that is to be saved as external data for callback functions.
+
+    Attributes:
+        total: The total number of tensors to save.
+        index: The index of the tensor being saved.
+        offset: The offset of the tensor in the external data file.
+    """
+
+    total: int
+    index: int
+    offset: int
+
+
 def _all_tensors(
     graph: _core.Graph | _core.GraphView, include_attributes: bool = False
 ) -> Iterator[_protocols.TensorProtocol]:
@@ -157,6 +175,7 @@ def _write_external_data(
     tensors: Sequence[_protocols.TensorProtocol],
     external_data_infos: Sequence[_ExternalDataInfo],
     file_path: str | os.PathLike,
+    callback: Callable[[_protocols.TensorProtocol, CallbackInfo], None] | None = None,
 ) -> None:
     """Write tensor data to an external file according to information stored in ExternalDataInfo objects.
 
@@ -164,12 +183,26 @@ def _write_external_data(
         tensors: Tensors to be written as external data.
         external_data_infos: External data information stored for each tensor to be written as external data.
         file_path: Location to which external data is to be stored.
+        callback: A callback function that is called for each tensor that is saved to external data
+            for debugging or logging purposes.
     """
-    assert len(tensors) == len(external_data_infos), (
+    tensors_count = len(tensors)
+    assert tensors_count == len(external_data_infos), (
         "Number of tensors and external data infos should match"
     )
     with open(file_path, "wb") as data_file:
-        for tensor, tensor_info in zip(tensors, external_data_infos, strict=True):
+        for i, (tensor, tensor_info) in enumerate(
+            zip(tensors, external_data_infos, strict=True)
+        ):
+            if callback is not None:
+                callback(
+                    tensor,
+                    CallbackInfo(
+                        total=tensors_count,
+                        index=i,
+                        offset=tensor_info.offset,
+                    ),
+                )
             current_offset = tensor_info.offset
             assert tensor is not None
             raw_data = tensor.tobytes()
@@ -228,6 +261,7 @@ def convert_tensors_to_external(
     tensors: Sequence[_protocols.TensorProtocol],
     base_dir: str | os.PathLike,
     relative_path: str | os.PathLike,
+    callback: Callable[[_protocols.TensorProtocol, CallbackInfo], None] | None = None,
 ) -> list[_core.ExternalTensor]:
     """Convert a sequence of any TensorProtocol tensors to external tensors.
 
@@ -238,6 +272,8 @@ def convert_tensors_to_external(
         tensors: Tensors to be converted to external tensors. They can be external tensors themselves.
         base_dir: Path of base directory.
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
+        callback: A callback function that is called for each tensor that is saved to external data
+            for debugging or logging purposes.
 
     Returns:
         A list of external tensors derived from a list of input tensors. The order
@@ -285,7 +321,7 @@ def convert_tensors_to_external(
         external_info = _compute_external_data_info(tensor, current_offset)
         external_data_infos.append(external_info)
         current_offset = external_info.offset + external_info.length
-    _write_external_data(sorted_tensors, external_data_infos, path)
+    _write_external_data(sorted_tensors, external_data_infos, path, callback=callback)
 
     # Create external tensor objects
     external_tensors: list[_core.ExternalTensor] = [
@@ -336,6 +372,7 @@ def unload_from_model(
     relative_path: str | os.PathLike,
     *,
     size_threshold_bytes: int = 0,
+    callback: Callable[[_protocols.TensorProtocol, CallbackInfo], None] | None = None,
 ) -> _core.Model:
     """Convert all initializers equal or above size_threshold_bytes to external tensors in-place and save data to a single data file.
 
@@ -356,6 +393,8 @@ def unload_from_model(
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
             E.g. "model.data"
         size_threshold_bytes: Save to external data if the tensor size in bytes is larger than this threshold.
+        callback: A callback function that is called for each tensor that is saved to external data
+            for debugging or logging purposes.
 
     Returns:
         An ir.Model with all initializer data equal or above ``size_threshold_bytes``
@@ -384,6 +423,7 @@ def unload_from_model(
         [v.const_value for v in initializers_to_become_external],  # type: ignore[misc]
         base_dir=base_dir,
         relative_path=relative_path,
+        callback=callback,
     )
 
     # Replace the initializer values with external tensors and save the model

onnx_ir/passes/_pass_infra.py

@@ -127,7 +127,7 @@ class PassBase(abc.ABC):
 
         # Check postconditions
         try:
-            self.ensures(model)
+            self.ensures(result.model)
         except PostconditionError:
             raise
         except Exception as e:

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/passes/common/_c_api_utils.py

@@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Callable, TypeVar
 import onnx_ir as ir
 
 if TYPE_CHECKING:
-    import onnx
+    import onnx  # noqa: TID251
 
 
 logger = logging.getLogger(__name__)

onnx_ir/passes/common/common_subexpression_elimination.py

@@ -17,93 +17,122 @@ logger = logging.getLogger(__name__)
 
 
 class CommonSubexpressionEliminationPass(ir.passes.InPlacePass):
-    """Eliminate common subexpression in ONNX graphs."""
+    """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."""
-        modified = False
         graph = model.graph
-
-        modified = _eliminate_common_subexpression(graph, modified)
+        modified = self._eliminate_common_subexpression(graph)
 
         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
+    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)
-            # 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,
+                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())),
             )
-            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
+            # 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(

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/passes/common/onnx_checker.py

@@ -10,7 +10,7 @@ __all__ = [
 
 from typing import Literal
 
-import onnx
+import onnx  # noqa: TID251
 
 import onnx_ir as ir
 from onnx_ir.passes.common import _c_api_utils

onnx_ir/passes/common/shape_inference.py

@@ -11,7 +11,7 @@ __all__ = [
 
 import logging
 
-import onnx
+import onnx  # noqa: TID251
 
 import onnx_ir as ir
 from onnx_ir.passes.common import _c_api_utils

onnx_ir/passes/common/unused_removal.py

@@ -10,7 +10,7 @@ __all__ = [
 
 import logging
 
-import onnx
+import onnx  # noqa: TID251
 
 import onnx_ir as ir