onnx-ir 0.1.11__tar.gz → 0.1.12__tar.gz

{onnx_ir-0.1.11/src/onnx_ir.egg-info → onnx_ir-0.1.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-ir
-Version: 0.1.11
+Version: 0.1.12
 Summary: Efficient in-memory representation for ONNX
 Author-email: ONNX Contributors <onnx-technical-discuss@lists.lfaidata.foundation>
 License-Expression: Apache-2.0
@@ -20,7 +20,6 @@ Dynamic: license-file
 # <img src="docs/_static/logo-light.png" alt="ONNX IR" width="250"/>
 
 [![PyPI - Version](https://img.shields.io/pypi/v/onnx-ir.svg)](https://pypi.org/project/onnx-ir)
-[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/onnx-ir.svg)](https://pypi.org/project/onnx-ir)
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![codecov](https://codecov.io/gh/onnx/ir-py/graph/badge.svg?token=SPQ3G9T78Z)](https://codecov.io/gh/onnx/ir-py)
 [![PyPI Downloads](https://static.pepy.tech/badge/onnx-ir/month)](https://pepy.tech/projects/onnx-ir)

{onnx_ir-0.1.11 → onnx_ir-0.1.12}/README.md

@@ -1,7 +1,6 @@
 # <img src="docs/_static/logo-light.png" alt="ONNX IR" width="250"/>
 
 [![PyPI - Version](https://img.shields.io/pypi/v/onnx-ir.svg)](https://pypi.org/project/onnx-ir)
-[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/onnx-ir.svg)](https://pypi.org/project/onnx-ir)
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![codecov](https://codecov.io/gh/onnx/ir-py/graph/badge.svg?token=SPQ3G9T78Z)](https://codecov.io/gh/onnx/ir-py)
 [![PyPI Downloads](https://static.pepy.tech/badge/onnx-ir/month)](https://pepy.tech/projects/onnx-ir)

{onnx_ir-0.1.11 → onnx_ir-0.1.12}/src/onnx_ir/__init__.py

@@ -168,4 +168,4 @@ def __set_module() -> None:
 
 
 __set_module()
-__version__ = "0.1.11"
+__version__ = "0.1.12"

{onnx_ir-0.1.11 → onnx_ir-0.1.12}/src/onnx_ir/_convenience/__init__.py

@@ -280,6 +280,7 @@ def convert_attributes(
 def replace_all_uses_with(
     values: _protocols.ValueProtocol | Sequence[_protocols.ValueProtocol],
     replacements: _protocols.ValueProtocol | Sequence[_protocols.ValueProtocol],
+    replace_graph_outputs: bool = False,
 ) -> None:
     """Replace all uses of the given values with the replacements.
 
@@ -318,9 +319,22 @@ def replace_all_uses_with(
         replaced are part of the graph outputs. Be sure to remove the old nodes
         from the graph using ``graph.remove()`` if they are no longer needed.
 
+    .. versionadded:: 0.1.12
+        The ``replace_graph_outputs`` parameter is added.
+
+    .. versionadded:: 0.1.12
+        ValueError is raised when ``replace_graph_outputs`` is False && when the value to
+        replace is a graph output.
+
     Args:
         values: The value or values to be replaced.
         replacements: The new value or values to use as inputs.
+        replace_graph_outputs: If True, graph outputs that reference the values
+            being replaced will also be updated to reference the replacements.
+
+    Raises:
+        ValueError: When ``replace_graph_outputs`` is False && when the value to
+            replace is a graph output.
     """
     if not isinstance(values, Sequence):
         values = (values,)
@@ -329,8 +343,7 @@ def replace_all_uses_with(
     if len(values) != len(replacements):
         raise ValueError("The number of values and replacements must match.")
     for value, replacement in zip(values, replacements):
-        for user_node, index in tuple(value.uses()):
-            user_node.replace_input_with(index, replacement)
+        value.replace_all_uses_with(replacement, replace_graph_outputs=replace_graph_outputs)
 
 
 def create_value_mapping(graph: _core.Graph) -> dict[str, _core.Value]:
@@ -397,20 +410,18 @@ def replace_nodes_and_values(
     """
     for old_value, new_value in zip(old_values, new_values):
         # Propagate relevant info from old value to new value
-        # TODO(Rama): Perhaps this should be a separate utility function. Also, consider
-        # merging old and new type/shape info.
-        new_value.type = old_value.type
-        new_value.shape = old_value.shape
-        new_value.const_value = old_value.const_value
-        new_value.name = old_value.name
+        # TODO(Rama): Perhaps this should be a separate utility function.
+        new_value.type = old_value.type if old_value.type is not None else new_value.type
+        new_value.shape = old_value.shape if old_value.shape is not None else new_value.shape
+        new_value.const_value = (
+            old_value.const_value
+            if old_value.const_value is not None
+            else new_value.const_value
+        )
+        new_value.name = old_value.name if old_value.name is not None else new_value.name
 
     # Reconnect the users of the deleted values to use the new values
-    replace_all_uses_with(old_values, new_values)
-    # Update graph/function outputs if the node generates output
-    replacement_mapping = dict(zip(old_values, new_values))
-    for idx, graph_or_function_output in enumerate(graph_or_function.outputs):
-        if graph_or_function_output in replacement_mapping:
-            graph_or_function.outputs[idx] = replacement_mapping[graph_or_function_output]
+    replace_all_uses_with(old_values, new_values, replace_graph_outputs=True)
 
     # insert new nodes after the index node
     graph_or_function.insert_after(insertion_point, new_nodes)

{onnx_ir-0.1.11 → onnx_ir-0.1.12}/src/onnx_ir/_core.py

@@ -2448,6 +2448,55 @@ class Value(_protocols.ValueProtocol, _display.PrettyPrintable):
         """Whether the value is an initializer of a graph."""
         return self._is_initializer
 
+    def replace_all_uses_with(
+        self, replacement: Value, /, replace_graph_outputs: bool = False
+    ) -> None:
+        """Replace all uses of this value with another value.
+
+        If the value is an output of a graph and ``replace_graph_outputs`` is ``True``,
+        the graph output will also be replaced. Be careful when a value appears multiple times
+        in the graph outputs - this is invalid. An identity node will need to be added on each
+        duplicated outputs to ensure a valid ONNX graph.
+
+        You may also want to assign the name of this value to the replacement value
+        to maintain the name when it is a graph output.
+
+        To replace usage of a sequence of values with another sequence of values, consider using
+        :func:`onnx_ir.convenience.replace_all_uses_with`.
+
+        .. versionadded:: 0.1.12
+
+        Args:
+            replacement: The value to replace all uses with.
+            replace_graph_outputs: If True, graph outputs that reference this value
+                will also be updated to reference the replacement.
+
+        Raises:
+            ValueError: When ``replace_graph_outputs`` is False && when the value to
+                replace is a graph output.
+        """
+        # NOTE: Why we don't replace the value name when the value is an output:
+        # When the replacement value is already an output of the graph, renaming it
+        # to the name of this value will cause name conflicts. It is better to let
+        # the user handle the renaming explicitly and insert identity nodes if needed.
+        if self.is_graph_output():
+            graph = self.graph
+            assert graph is not None
+
+            if not replace_graph_outputs:
+                raise ValueError(
+                    f"{self!r} is an output of graph {graph.name!r}. "
+                    "Set replace_graph_outputs=True or replace the graph output frist before "
+                    "calling replace_all_uses_with."
+                )
+
+            for i, output in enumerate(graph.outputs):
+                if output is self:
+                    graph.outputs[i] = replacement
+
+        for user_node, index in self.uses():
+            user_node.replace_input_with(index, replacement)
+
 
 @deprecated("Input is deprecated since 0.1.9. Use ir.val(...) instead.")
 def Input(

{onnx_ir-0.1.11 → onnx_ir-0.1.12}/src/onnx_ir/_protocols.py

@@ -203,6 +203,17 @@ class ValueProtocol(Protocol):
         """Whether this value is an output of a graph."""
         ...
 
+    def replace_all_uses_with(
+        self, new_value: ValueProtocol | None, replace_graph_outputs: bool = False
+    ) -> None:
+        """Replace all uses of this value with the given new value.
+
+        Args:
+            new_value: The new value to replace this value with.
+            replace_graph_outputs: Whether to replace graph outputs that use this value.
+        """
+        ...
+
 
 @typing.runtime_checkable
 class NodeProtocol(Protocol):

{onnx_ir-0.1.11 → onnx_ir-0.1.12}/src/onnx_ir/passes/common/common_subexpression_elimination.py

@@ -150,8 +150,6 @@ def _remove_node_and_replace_values(
         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))
@@ -185,6 +183,9 @@ def _remove_node_and_replace_values(
                     new_value.name = graph_output.name
                     graph.outputs[idx] = new_value
 
+    # Reconnect the users of the deleted values to use the new values
+    ir.convenience.replace_all_uses_with(remove_values, new_values)
+
     graph.remove(remove_node, safe=True)
 
 

{onnx_ir-0.1.11 → onnx_ir-0.1.12}/src/onnx_ir/passes/common/constant_manipulation.py

@@ -78,7 +78,7 @@ class LiftConstantsToInitializersPass(ir.passes.InPlacePass):
             assert node.graph is not None
             node.graph.register_initializer(initializer)
             # Replace the constant node with the initializer
-            ir.convenience.replace_all_uses_with(node.outputs[0], initializer)
+            node.outputs[0].replace_all_uses_with(initializer)
             node.graph.remove(node, safe=True)
             count += 1
             logger.debug(

{onnx_ir-0.1.11 → onnx_ir-0.1.12}/src/onnx_ir/passes/common/identity_elimination.py

@@ -105,20 +105,14 @@ class IdentityEliminationPass(ir.passes.InPlacePass):
 
         # Case 1 & 2 (merged): Eliminate the identity node
         # Replace all uses of output with input
-        ir.convenience.replace_all_uses_with(output_value, input_value)
+        ir.convenience.replace_all_uses_with(
+            output_value, input_value, replace_graph_outputs=True
+        )
 
         # If output is a graph output, we need to rename input and update graph outputs
         if output_is_graph_output:
-            # Store the original output name
-            original_output_name = output_value.name
-
             # Update the input value to have the output's name
-            input_value.name = original_output_name
-
-            # Update graph outputs to point to the input value
-            for idx, graph_output in enumerate(graph_like.outputs):
-                if graph_output is output_value:
-                    graph_like.outputs[idx] = input_value
+            input_value.name = output_value.name
 
         # Remove the identity node
         graph_like.remove(node, safe=True)

{onnx_ir-0.1.11 → onnx_ir-0.1.12}/src/onnx_ir/passes/common/initializer_deduplication.py

@@ -100,7 +100,7 @@ class DeduplicateInitializersPass(ir.passes.InPlacePass):
                 if key in initializers:
                     modified = True
                     initializer_to_keep = initializers[key]  # type: ignore[index]
-                    ir.convenience.replace_all_uses_with(initializer, initializer_to_keep)
+                    initializer.replace_all_uses_with(initializer_to_keep)
                     assert initializer.name is not None
                     graph.initializers.pop(initializer.name)
                     logger.info(
@@ -165,7 +165,7 @@ class DeduplicateHashedInitializersPass(ir.passes.InPlacePass):
                         continue
                     modified = True
                     initializer_to_keep = initializers[key]  # type: ignore[index]
-                    ir.convenience.replace_all_uses_with(initializer, initializer_to_keep)
+                    initializer.replace_all_uses_with(initializer_to_keep)
                     assert initializer.name is not None
                     graph.initializers.pop(initializer.name)
                     logger.info(

{onnx_ir-0.1.11 → onnx_ir-0.1.12/src/onnx_ir.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-ir
-Version: 0.1.11
+Version: 0.1.12
 Summary: Efficient in-memory representation for ONNX
 Author-email: ONNX Contributors <onnx-technical-discuss@lists.lfaidata.foundation>
 License-Expression: Apache-2.0
@@ -20,7 +20,6 @@ Dynamic: license-file
 # <img src="docs/_static/logo-light.png" alt="ONNX IR" width="250"/>
 
 [![PyPI - Version](https://img.shields.io/pypi/v/onnx-ir.svg)](https://pypi.org/project/onnx-ir)
-[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/onnx-ir.svg)](https://pypi.org/project/onnx-ir)
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![codecov](https://codecov.io/gh/onnx/ir-py/graph/badge.svg?token=SPQ3G9T78Z)](https://codecov.io/gh/onnx/ir-py)
 [![PyPI Downloads](https://static.pepy.tech/badge/onnx-ir/month)](https://pepy.tech/projects/onnx-ir)