onnx-ir 0.1.11__py3-none-any.whl → 0.1.12__py3-none-any.whl

onnx_ir/__init__.py

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

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/_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/_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/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/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/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/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.dist-info → onnx_ir-0.1.12.dist-info}/METADATA

@@ -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.dist-info → onnx_ir-0.1.12.dist-info}/RECORD

@@ -1,5 +1,5 @@
-onnx_ir/__init__.py,sha256=152XzM3zz66eHZYCpERRhPZhEExuXkOhdS_zYi_0wQE,3441
-onnx_ir/_core.py,sha256=32tCB7gZkRjI2kcrp5cLUGz-KH7-BAbaxVeuoWR2TZs,147035
+onnx_ir/__init__.py,sha256=bPBQNMnNPrPCuJslVrQdyzwxDjsjPyWqqtR7rOt95vM,3441
+onnx_ir/_core.py,sha256=A5Zz7AadmKCX_ibMeP7APJUa06xZBL1poj8KBv9N1IE,149270
 onnx_ir/_display.py,sha256=230bMN_hVy47Ug3HkA4o5Tf5Hr21AnBEoq5w0fxjyTs,1300
 onnx_ir/_enums.py,sha256=E7WQ7yQzulBeimamc9q_k4fEUoyH_2PWtaOMpwck_W0,13915
 onnx_ir/_graph_comparison.py,sha256=8_D1gu547eCDotEUqxfIJhUGU_Ufhfji7sfsSraOj3g,727
@@ -9,7 +9,7 @@ onnx_ir/_linked_list.py,sha256=PXVcbHLMXHLZ6DxZnElnJLWfhBPvYcXUxM8Y3K4J7lM,10592
 onnx_ir/_metadata.py,sha256=lzmCaYy4kAJrPW-PSGKF4a78LisxF0hiofySNQ3Mwhg,1544
 onnx_ir/_name_authority.py,sha256=PnoV9TRgMLussZNufWavJXosDWx5avPfldVjMWEEz18,3036
 onnx_ir/_polyfill.py,sha256=LzAGBKQbVDlURC0tgQgaxgkYU4rESgCYnqVs-u-Vsx8,887
-onnx_ir/_protocols.py,sha256=M29sIOAvtdlis3QtBvCQPH4pnvSwhJCQNCvs3IrN9FY,21276
+onnx_ir/_protocols.py,sha256=PHJtdhATDNwnfocIpvQUNM2dPn1sxfGbU1JPhFDvoIA,21667
 onnx_ir/_tape.py,sha256=nEGY6VZVKuB8FDyXeYr0MTq8j7E4HKOE2yN8qpz4ia0,7007
 onnx_ir/_type_casting.py,sha256=hbikTmgFEu0SEfnbgv2R1LbpuPQ2MCfqto3-oLWhcBc,1645
 onnx_ir/_version_utils.py,sha256=bZThuE7meVHFOY1DLsmss9WshVIp9iig7udGfDbVaK4,1333
@@ -21,7 +21,7 @@ onnx_ir/tape.py,sha256=4FyfAHmVhQoMsfHMYnBwP2azi6UF6b6pj--ercObqZs,350
 onnx_ir/tensor_adapters.py,sha256=1YxOgYro9QV-Dw4XsiD9WQgs1_6h13nu6tfmA1mG7IA,7568
 onnx_ir/testing.py,sha256=WTrjf2joWizDWaYMJlV1KjZMQw7YmZ8NvuBTVn1uY6s,8803
 onnx_ir/traversal.py,sha256=Wy4XphwuapAvm94-5iaz6G8LjIoMFpY7qfPfXzYViEE,4488
-onnx_ir/_convenience/__init__.py,sha256=SO7kc8RXVKEUODGh0q2Y7WgmbUsOjYSixmKFx_A0DAQ,19752
+onnx_ir/_convenience/__init__.py,sha256=UWaFlqLwgdk09Mn-9O9jk-ZZHiBqvDzP_zFpMgNp5fw,20183
 onnx_ir/_convenience/_constructors.py,sha256=HqCGtNPMzFFEerC7I5VEyMdBuIdOJDucn9UEdwuymcg,11519
 onnx_ir/_thirdparty/asciichartpy.py,sha256=afQ0fsqko2uYRPAR4TZBrQxvCb4eN8lxZ2yDFbVQq_s,10533
 onnx_ir/passes/__init__.py,sha256=M_Tcl_-qGSNPluFIvOoeDyh0qAwNayaYyXDS5UJUJPQ,764
@@ -29,18 +29,18 @@ onnx_ir/passes/_pass_infra.py,sha256=xIOw_zZIuOqD4Z_wZ4OvsqXfh2IZMoMlDp1xQ_MPQlc
 onnx_ir/passes/common/__init__.py,sha256=NYBrXvkq_CbWRwcZptSsF4u_-1zfN_BClLhVQY0pwYc,1738
 onnx_ir/passes/common/_c_api_utils.py,sha256=g6riA6xNGVWaO5YjVHZ0krrfslWHmRlryRkwB8X56cg,2907
 onnx_ir/passes/common/clear_metadata_and_docstring.py,sha256=YwouLfsNFSaTuGd7uMOGjdvVwG9yHQTkSphUgDlM0ME,2365
-onnx_ir/passes/common/common_subexpression_elimination.py,sha256=wZ1zEPdCshYB_ifP9fCAVfzQkesE6uhCfzCuL2qO5fA,7948
-onnx_ir/passes/common/constant_manipulation.py,sha256=Ja_2uO59ni8YX600iF3wvCmOk4EvM5crl7csYU1s7rQ,9391
-onnx_ir/passes/common/identity_elimination.py,sha256=yLoy0DEIIfhP1smlCx-2Bxsv89lij1uvFdFVh9XUfTY,4667
-onnx_ir/passes/common/initializer_deduplication.py,sha256=gKrXTMFAtCkMmiIm8zWzwPnwSbRdZxunJeAt_jFU-vY,7253
+onnx_ir/passes/common/common_subexpression_elimination.py,sha256=p5hZjWyswn8qm91M6BqlqMUBu78ohxaw72Ky12zrPZ0,7949
+onnx_ir/passes/common/constant_manipulation.py,sha256=qKC--Mr8IzrXPmr6xKJKOVO54M9wh8Ruho423LQ1mfU,9375
+onnx_ir/passes/common/identity_elimination.py,sha256=90HYaq4QYUilHMyRoN_qhwWt0bQpqm5Z69KNU_7z_As,4376
+onnx_ir/passes/common/initializer_deduplication.py,sha256=2OK6h6cLp2VmfT5VUxsknsRXr1fgbY2w5npn0hV1cdE,7221
 onnx_ir/passes/common/inliner.py,sha256=wBoO6yXt6F1AObQjYZHMQ0wn3YH681N4HQQVyaMAYd4,13702
 onnx_ir/passes/common/naming.py,sha256=l_LrUiI3gAoSEVs8YeQ5kRNWp7aMOlBK-SfFUsKobZI,10687
 onnx_ir/passes/common/onnx_checker.py,sha256=_sPmJ2ff9pDB1g9q7082BL6fyubomRaj6svE0cCyDew,1691
 onnx_ir/passes/common/shape_inference.py,sha256=LVdvxjeKtcIEbPcb6mKisxoPJOOawzsm3tzk5j9xqeM,3992
 onnx_ir/passes/common/topological_sort.py,sha256=Vcu1YhBdfRX4LROr0NScjB1Pwz2DjBFD0Z_GxqaxPF8,999
 onnx_ir/passes/common/unused_removal.py,sha256=cBNqaqGnUVyCWxsD7hBzYk4qSglVPo3SmHAvkUo5-Oc,7613
-onnx_ir-0.1.11.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-onnx_ir-0.1.11.dist-info/METADATA,sha256=Q5F7zhOGZspqkN7nI0NnAjSdmRgZqFK6zgWM-2azglA,3351
-onnx_ir-0.1.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-onnx_ir-0.1.11.dist-info/top_level.txt,sha256=W5tROO93YjO0XRxIdjMy4wocp-5st5GiI2ukvW7UhDo,8
-onnx_ir-0.1.11.dist-info/RECORD,,
+onnx_ir-0.1.12.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+onnx_ir-0.1.12.dist-info/METADATA,sha256=cys4n_vm2TolOB7h4oGA-LPc0ippsJKJf4hzqoa7bXM,3238
+onnx_ir-0.1.12.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+onnx_ir-0.1.12.dist-info/top_level.txt,sha256=W5tROO93YjO0XRxIdjMy4wocp-5st5GiI2ukvW7UhDo,8
+onnx_ir-0.1.12.dist-info/RECORD,,