onnx-ir 0.0.1__py3-none-any.whl

onnx_ir/traversal.py

@@ -0,0 +1,82 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Utilities for traversing the IR graph."""
+
+from __future__ import annotations
+
+__all__ = [
+    "RecursiveGraphIterator",
+]
+
+from typing import Callable, Iterator, Reversible, Union
+
+from typing_extensions import Self
+
+from onnx_ir import _core, _enums
+
+GraphLike = Union[_core.Graph, _core.Function, _core.GraphView]
+
+
+class RecursiveGraphIterator(Iterator[_core.Node], Reversible[_core.Node]):
+    def __init__(
+        self,
+        graph_like: GraphLike,
+        *,
+        recursive: Callable[[_core.Node], bool] | None = None,
+        reverse: bool = False,
+    ):
+        """Iterate over the nodes in the graph, recursively visiting subgraphs.
+
+        Args:
+            graph_like: The graph to traverse.
+            recursive: A callback that determines whether to recursively visit the subgraphs
+                contained in a node. If not provided, all nodes in subgraphs are visited.
+            reverse: Whether to iterate in reverse order.
+        """
+        self._graph = graph_like
+        self._recursive = recursive
+        self._reverse = reverse
+        self._iterator = self._recursive_node_iter(graph_like)
+
+    def __iter__(self) -> Self:
+        self._iterator = self._recursive_node_iter(self._graph)
+        return self
+
+    def __next__(self) -> _core.Node:
+        return next(self._iterator)
+
+    def _recursive_node_iter(
+        self, graph: _core.Graph | _core.Function | _core.GraphView
+    ) -> Iterator[_core.Node]:
+        iterable = reversed(graph) if self._reverse else graph
+        for node in iterable:  # type: ignore[union-attr]
+            yield node
+            if self._recursive is not None and not self._recursive(node):
+                continue
+            yield from self._iterate_subgraphs(node)
+
+    def _iterate_subgraphs(self, node: _core.Node):
+        for attr in node.attributes.values():
+            if not isinstance(attr, _core.Attr):
+                continue
+            if attr.type == _enums.AttributeType.GRAPH:
+                yield from RecursiveGraphIterator(
+                    attr.value,
+                    recursive=self._recursive,
+                    reverse=self._reverse,
+                )
+            elif attr.type == _enums.AttributeType.GRAPHS:
+                graphs = reversed(attr.value) if self._reverse else attr.value
+                for graph in graphs:
+                    yield from RecursiveGraphIterator(
+                        graph,
+                        recursive=self._recursive,
+                        reverse=self._reverse,
+                    )
+
+    def __reversed__(self) -> Iterator[_core.Node]:
+        return RecursiveGraphIterator(
+            self._graph,
+            recursive=self._recursive,
+            reverse=not self._reverse,
+        )

onnx_ir-0.0.1.dist-info/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2025 Justin Chu
+Copyright (c) Microsoft Corporation
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

onnx_ir-0.0.1.dist-info/METADATA

@@ -0,0 +1,73 @@
+Metadata-Version: 2.2
+Name: onnx-ir
+Version: 0.0.1
+Author-email: Justin Chu <justinchu@microsoft.com>
+License: MIT License
+        
+        Copyright (c) 2025 Justin Chu
+        Copyright (c) Microsoft Corporation
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Classifier: Development Status :: 4 - Beta
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: POSIX
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: License :: OSI Approved :: MIT License
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: onnx>=1.16
+Requires-Dist: typing_extensions>=4.10
+Requires-Dist: ml_dtypes
+
+# ONNX IR
+
+> [!NOTE]
+> This is an adaptation of the `onnxscript.ir` module from the ONNX Script project: https://github.com/microsoft/onnxscript/tree/main/onnxscript/ir
+
+An in-memory IR that supports the full ONNX spec, designed for graph construction, analysis and transformation.
+
+## 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).
+- Low memory footprint: mmap'ed external tensors; unified interface for ONNX TensorProto, Numpy arrays and PyTorch Tensors etc. No tensor size limitation. Zero copies.
+- Straightforward access patterns: Access value information and traverse the graph topology at ease.
+- Robust mutation: Create as many iterators as you like on the graph while mutating it.
+- Speed: Performant graph manipulation, serialization/deserialization to Protobuf.
+- Pythonic and familiar APIs: Classes define Pythonic apis and still map to ONNX protobuf concepts in an intuitive way.
+- No protobuf dependency: The IR does not require protobuf once the model is converted to the IR representation, decoupling from the serialization format.
+
+## Code Organization 🗺️
+
+- [`_protocols.py`](_protocols.py): Interfaces defined for all entities in the IR.
+- [`_core.py`](_core.py): Implementation of the core entities in the IR, including `Model`, `Graph`, `Node`, `Value`, and others.
+- [`_enums.py`](_enums.py): Definition of the type enums that correspond to the `DataType` and `AttributeType` in `onnx.proto`.
+- [`_name_authority.py`](_name_authority.py): The authority for giving names to entities in the graph, used internally.
+- [`_linked_list.py`](_linked_list.py): The data structure as the node container in the graph that supports robust iteration and mutation. Internal.
+- [`_metadata.py`](_metadata.py): Metadata store for all entities in the IR.

onnx_ir-0.0.1.dist-info/RECORD

@@ -0,0 +1,26 @@
+onnx_ir/__init__.py,sha256=Lj7uEV33mcR9dLgvA0Qs-pFkE3faJoLZLtJGC-QphjM,2993
+onnx_ir/_convenience.py,sha256=G7uDrvbKYpk65-f-wYqORKuP0roxMQ-miMiGEMnHDG8,16467
+onnx_ir/_core.py,sha256=0TS4bzBvsChq-uDueFLMU32azRIxY8HDF9oIXRCnDGw,103991
+onnx_ir/_display.py,sha256=6x5QQwpA4ugEhbCZ8ZzFGMNcpNlnjXQIftJrT-cb61c,1293
+onnx_ir/_enums.py,sha256=0Q-1ee0XzM0EDwX-CbfS_i48XjLErtwDVz7H34vDUFg,4127
+onnx_ir/_external_data.py,sha256=xV4Fi_jBXcbvMRzfbdeFH5rZMAi2bym2E9NdCqtABQM,13747
+onnx_ir/_graph_comparison.py,sha256=wGkD4HLfC4JOcvOSxNM1dLMpc1QPEy_uCnasL5Z_2RY,720
+onnx_ir/_io.py,sha256=Nnid_zzuA7G0a5IhJHvI6smLqYtftXpC6ILKtHpWhp8,1703
+onnx_ir/_linked_list.py,sha256=rTog1WP4YFfXkOtvc63X7zzKvFydneVn3I3W2vv9JEc,10344
+onnx_ir/_metadata.py,sha256=4SL4m-GiRO3z0fcOfXzye-5QwFwKxFacqoVx3_AzyVI,1510
+onnx_ir/_name_authority.py,sha256=B7fS92M1zTz0iSwTbxzYt1fbaOuQ2kv7-wRZC8zkqoQ,3029
+onnx_ir/_protocols.py,sha256=PgrWncyCERua4fMVnmwPJnfD08i-9y8dYCNdqY2fA3I,20818
+onnx_ir/_tape.py,sha256=bRg7LipHIZr62TbLbAB-rjjomDDkwoupDuWj27_vGv0,3417
+onnx_ir/_type_casting.py,sha256=K8W6NEwbsQWQ-BSl9j928-vz6Posy0mUnMGBjX2ZAj0,2848
+onnx_ir/convenience.py,sha256=QRVmuU8OSS4mu-VOjv0MjHHxpQC5VOKAGnwHjWmvOJ4,810
+onnx_ir/serde.py,sha256=YXJD7146RERbPK5SQBJe-lYlhgs-0Ir71hmw7SdZc1g,62852
+onnx_ir/traversal.py,sha256=8zw8XvxUefzbWfGcK28c9h4XTnc7IqNgKYCZnIB0Q-U,2829
+onnx_ir/_internal/version_utils.py,sha256=fDJMHm5lxv-Q1Iz4KJ8kM9pmFFUymz1zi0a_uDF0Js4,3281
+onnx_ir/_thirdparty/asciichartpy.py,sha256=wgqaJBjSZRvftSzQO8TPyldnYt5wUlApKAc9GxFVZpw,10599
+onnx_ir/passes/__init__.py,sha256=O2uRd8pbRhYYMzlKJ4CnWNq3MR2JXzQCv_lFFIZHc4A,645
+onnx_ir/passes/_pass_infra.py,sha256=ygI2Ptq5Moleq2pr23OPnVmQVY4e0LyupQY0c6-BRwM,5366
+onnx_ir-0.0.1.dist-info/LICENSE,sha256=mwKiiaO6cDOyDz3J5_ApIDbydWOxYwi8jD2WOVHrJX4,1103
+onnx_ir-0.0.1.dist-info/METADATA,sha256=6IJ-7WF9AB4Sl3WHRV9LnPz74MY-9ylIJB-kIB0LQYQ,4076
+onnx_ir-0.0.1.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+onnx_ir-0.0.1.dist-info/top_level.txt,sha256=W5tROO93YjO0XRxIdjMy4wocp-5st5GiI2ukvW7UhDo,8
+onnx_ir-0.0.1.dist-info/RECORD,,

onnx_ir-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.8.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

onnx_ir-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+onnx_ir