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

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/_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/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
 

onnx_ir/serde.py

@@ -37,6 +37,7 @@ __all__ = [
     "deserialize_value_info_proto",
     # Serialization
     "to_proto",
+    "to_onnx_text",
     "serialize_attribute_into",
     "serialize_attribute",
     "serialize_dimension_into",
@@ -62,14 +63,14 @@ __all__ = [
 import collections
 import logging
 import os
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Any, Callable
 
 import numpy as np
-import onnx
-import onnx.external_data_helper
+import onnx  # noqa: TID251
+import onnx.external_data_helper  # noqa: TID251
 
-from onnx_ir import _core, _enums, _protocols, _type_casting
+from onnx_ir import _convenience, _core, _enums, _protocols, _type_casting
 
 if typing.TYPE_CHECKING:
     import google.protobuf.internal.containers as proto_containers
@@ -190,13 +191,64 @@ def from_proto(proto: object) -> object:
     )
 
 
-def from_onnx_text(model_text: str, /) -> _core.Model:
+def from_onnx_text(
+    model_text: str,
+    /,
+    initializers: Iterable[_protocols.TensorProtocol] | None = None,
+) -> _core.Model:
     """Convert the ONNX textual representation to an IR model.
 
     Read more about the textual representation at: https://onnx.ai/onnx/repo-docs/Syntax.html
+
+    Args:
+        model_text: The ONNX textual representation of the model.
+        initializers: Tensors to be added as initializers. If provided, these tensors
+            will be added to the model as initializers. If a name does not exist in the model,
+            a ValueError will be raised.
+
+    Returns:
+        The IR model corresponding to the ONNX textual representation.
+
+    Raises:
+        ValueError: If a tensor name in `initializers` does not match any value in the model.
     """
     proto = onnx.parser.parse_model(model_text)
-    return deserialize_model(proto)
+    model = deserialize_model(proto)
+    values = _convenience.create_value_mapping(model.graph)
+    if initializers:
+        # Add initializers to the model
+        for tensor in initializers:
+            name = tensor.name
+            if not name:
+                raise ValueError(
+                    "Initializer tensor must have a name. "
+                    f"Please provide a name for the initializer: {tensor}"
+                )
+            if name not in values:
+                raise ValueError(f"Value '{name}' does not exist in model.")
+            initializer = values[name]
+            initializer.const_value = tensor
+            model.graph.register_initializer(initializer)
+    return model
+
+
+def to_onnx_text(
+    model: _protocols.ModelProtocol, /, exclude_initializers: bool = False
+) -> str:
+    """Convert the IR model to the ONNX textual representation.
+
+    Args:
+        model: The IR model to convert.
+        exclude_initializers: If True, the initializers will not be included in the output.
+
+    Returns:
+        The ONNX textual representation of the model.
+    """
+    proto = serialize_model(model)
+    if exclude_initializers:
+        del proto.graph.initializer[:]
+    text = onnx.printer.to_text(proto)
+    return text
 
 
 @typing.overload
@@ -462,6 +514,14 @@ def _get_field(proto: Any, field: str) -> Any:
 def deserialize_opset_import(
     protos: Sequence[onnx.OperatorSetIdProto],
 ) -> dict[str, int]:
+    """Deserialize a sequence of OperatorSetIdProto to opset imports mapping.
+
+    Args:
+        protos: The sequence of ONNX OperatorSetIdProto objects.
+
+    Returns:
+        A dictionary mapping domain strings to version integers.
+    """
     return {opset.domain: opset.version for opset in protos}
 
 
@@ -495,6 +555,14 @@ def _parse_experimental_function_value_info_name(
 
 
 def deserialize_model(proto: onnx.ModelProto) -> _core.Model:
+    """Deserialize an ONNX ModelProto into an IR Model.
+
+    Args:
+        proto: The ONNX ModelProto to deserialize.
+
+    Returns:
+        An IR Model object representing the ONNX model.
+    """
     graph = _deserialize_graph(proto.graph, [])
     graph.opset_imports.update(deserialize_opset_import(proto.opset_import))
 
@@ -699,6 +767,14 @@ def _deserialize_graph(
 
 @_capture_errors(lambda proto: proto.name)
 def deserialize_function(proto: onnx.FunctionProto) -> _core.Function:
+    """Deserialize an ONNX FunctionProto into an IR Function.
+
+    Args:
+        proto: The ONNX FunctionProto to deserialize.
+
+    Returns:
+        An IR Function object representing the ONNX function.
+    """
     inputs = [_core.Input(name) for name in proto.input]
     values: dict[str, _core.Value] = {v.name: v for v in inputs}  # type: ignore[misc]
     value_info = {info.name: info for info in getattr(proto, "value_info", [])}
@@ -741,6 +817,15 @@ def deserialize_function(proto: onnx.FunctionProto) -> _core.Function:
 def deserialize_value_info_proto(
     proto: onnx.ValueInfoProto, value: _core.Value | None
 ) -> _core.Value:
+    """Deserialize an ONNX ValueInfoProto into an IR Value.
+
+    Args:
+        proto: The ONNX ValueInfoProto to deserialize.
+        value: An existing Value to update, or None to create a new one.
+
+    Returns:
+        An IR Value object with type and shape information populated from the proto.
+    """
     if value is None:
         value = _core.Value(name=proto.name)
     value.shape = deserialize_type_proto_for_shape(proto.type)
@@ -767,6 +852,14 @@ def _deserialize_quantization_annotation(
 
 @_capture_errors(str)
 def deserialize_tensor_shape(proto: onnx.TensorShapeProto) -> _core.Shape:
+    """Deserialize an ONNX TensorShapeProto into an IR Shape.
+
+    Args:
+        proto: The ONNX TensorShapeProto to deserialize.
+
+    Returns:
+        An IR Shape object representing the tensor shape.
+    """
     # This logic handles when the shape is [] as well
     dim_protos = proto.dim
     deserialized_dim_denotations = [
@@ -779,6 +872,14 @@ def deserialize_tensor_shape(proto: onnx.TensorShapeProto) -> _core.Shape:
 
 @_capture_errors(str)
 def deserialize_type_proto_for_shape(proto: onnx.TypeProto) -> _core.Shape | None:
+    """Extract and deserialize shape information from an ONNX TypeProto.
+
+    Args:
+        proto: The ONNX TypeProto to extract shape from.
+
+    Returns:
+        An IR Shape object if shape information is present, None otherwise.
+    """
     if proto.HasField("tensor_type"):
         if (shape_proto := _get_field(proto.tensor_type, "shape")) is None:
             return None
@@ -806,6 +907,14 @@ def deserialize_type_proto_for_shape(proto: onnx.TypeProto) -> _core.Shape | Non
 def deserialize_type_proto_for_type(
     proto: onnx.TypeProto,
 ) -> _protocols.TypeProtocol | None:
+    """Extract and deserialize type information from an ONNX TypeProto.
+
+    Args:
+        proto: The ONNX TypeProto to extract type from.
+
+    Returns:
+        An IR type object (TensorType, SequenceType, etc.) if type information is present, None otherwise.
+    """
     denotation = _get_field(proto, "denotation")
     if proto.HasField("tensor_type"):
         if (elem_type := _get_field(proto.tensor_type, "elem_type")) is None:
@@ -906,6 +1015,14 @@ _deserialize_string_string_maps = deserialize_metadata_props
 
 
 def deserialize_attribute(proto: onnx.AttributeProto) -> _core.Attr:
+    """Deserialize an ONNX AttributeProto into an IR Attribute.
+
+    Args:
+        proto: The ONNX AttributeProto to deserialize.
+
+    Returns:
+        An IR Attribute object representing the ONNX attribute.
+    """
     return _deserialize_attribute(proto, [])
 
 
@@ -979,6 +1096,14 @@ def _deserialize_attribute(
 
 
 def deserialize_node(proto: onnx.NodeProto) -> _core.Node:
+    """Deserialize an ONNX NodeProto into an IR Node.
+
+    Args:
+        proto: The ONNX NodeProto to deserialize.
+
+    Returns:
+        An IR Node object representing the ONNX node.
+    """
     return _deserialize_node(
         proto, scoped_values=[{}], value_info={}, quantization_annotations={}
     )
@@ -1097,6 +1222,14 @@ def _deserialize_node(
 
 
 def serialize_model(model: _protocols.ModelProtocol) -> onnx.ModelProto:
+    """Serialize an IR Model to an ONNX ModelProto.
+
+    Args:
+        model: The IR Model to serialize.
+
+    Returns:
+        The serialized ONNX ModelProto object.
+    """
     return serialize_model_into(onnx.ModelProto(), from_=model)
 
 
@@ -1418,6 +1551,14 @@ def serialize_function_into(
 
 
 def serialize_node(node: _protocols.NodeProtocol) -> onnx.NodeProto:
+    """Serialize an IR Node to an ONNX NodeProto.
+
+    Args:
+        node: The IR Node to serialize.
+
+    Returns:
+        The serialized ONNX NodeProto object.
+    """
     node_proto = onnx.NodeProto()
     serialize_node_into(node_proto, from_=node)
     return node_proto
@@ -1472,6 +1613,14 @@ def serialize_node_into(node_proto: onnx.NodeProto, from_: _protocols.NodeProtoc
 
 
 def serialize_tensor(tensor: _protocols.TensorProtocol) -> onnx.TensorProto:
+    """Serialize an IR Tensor to an ONNX TensorProto.
+
+    Args:
+        tensor: The IR Tensor to serialize.
+
+    Returns:
+        The serialized ONNX TensorProto object.
+    """
     tensor_proto = onnx.TensorProto()
     serialize_tensor_into(tensor_proto, from_=tensor)
     return tensor_proto
@@ -1514,6 +1663,14 @@ def serialize_tensor_into(
 
 
 def serialize_attribute(attribute: _protocols.AttributeProtocol) -> onnx.AttributeProto:
+    """Serialize an IR Attribute to an ONNX AttributeProto.
+
+    Args:
+        attribute: The IR Attribute to serialize.
+
+    Returns:
+        The serialized ONNX AttributeProto object.
+    """
     attribute_proto = onnx.AttributeProto()
     serialize_attribute_into(attribute_proto, from_=attribute)
     return attribute_proto
@@ -1678,6 +1835,14 @@ def serialize_type_into(type_proto: onnx.TypeProto, from_: _protocols.TypeProtoc
 
 
 def serialize_type(type_protocol: _protocols.TypeProtocol) -> onnx.TypeProto:
+    """Serialize an IR Type to an ONNX TypeProto.
+
+    Args:
+        type_protocol: The IR Type to serialize.
+
+    Returns:
+        The serialized ONNX TypeProto object.
+    """
     type_proto = onnx.TypeProto()
     serialize_type_into(type_proto, from_=type_protocol)
     return type_proto

{onnx_ir-0.1.1.dist-info → onnx_ir-0.1.2.dist-info}/METADATA

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: onnx-ir
-Version: 0.1.1
+Version: 0.1.2
 Summary: Efficient in-memory representation for ONNX
 Author-email: ONNX Contributors <onnx-technical-discuss@lists.lfaidata.foundation>
 License: Apache License v2.0
-Project-URL: Homepage, https://onnx.ai/onnx-ir
-Project-URL: Issues, https://github.com/onnx/onnx-ir/issues
-Project-URL: Repository, https://github.com/onnx/onnx-ir
+Project-URL: Homepage, https://onnx.ai/ir-py
+Project-URL: Issues, https://github.com/onnx/ir-py/issues
+Project-URL: Repository, https://github.com/onnx/ir-py
 Classifier: Development Status :: 4 - Beta
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
@@ -33,6 +33,24 @@ Dynamic: license-file
 
 An in-memory IR that supports the full ONNX spec, designed for graph construction, analysis and transformation.
 
+## Getting Started
+
+[onnx-ir documentation](https://onnx.ai/ir-py/)
+
+### Installation
+
+Via pip:
+
+```
+pip install onnx-ir
+```
+
+Or from source:
+
+```
+pip install git+https://github.com/onnx/ir-py.git
+```
+
 ## Features ✨
 
 - Full ONNX spec support: all valid models representable by ONNX protobuf, and a subset of invalid models (so you can load and fix them).