Improve documentation on replace_all_uses_with (onnx#237)

Update the docstring to reflect current usage and improve tips.

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

src/onnx_ir/_convenience/__init__.py

@@ -315,9 +315,23 @@ def replace_all_uses_with(
     users of the first value is replaced with the first replacement, and so on.
 
     .. note::
-        You still need to update the graph outputs if any of the values being
-        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.
+        Be sure to remove the old nodes from the graph using ``graph.remove()``
+        if they are no longer needed, or use :class:`onnx_ir.passes.common.RemoveUnusedNodesPass`
+        to remove all unused nodes in the graph.
+
+    .. tip::
+        **Handling graph outputs**
+
+        To also replace graph outputs that reference the values being replaced, either
+        set ``replace_graph_outputs`` to True, or manually update the graph outputs
+        before calling this function to avoid an error being raised when ``replace_graph_outputs=False``.
+
+        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.
 
     .. versionadded:: 0.1.12
         The ``replace_graph_outputs`` parameter is added.

src/onnx_ir/_core.py

@@ -2453,13 +2453,19 @@ def replace_all_uses_with(
     ) -> 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.
+        .. tip::
+            **Handling graph outputs**
 
-        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 also replace graph outputs that reference the values being replaced, either
+            set ``replace_graph_outputs`` to True, or manually update the graph outputs
+            before calling this function to avoid an error being raised when ``replace_graph_outputs=False``.
+
+            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`.

src/onnx_ir/passes/common/topological_sort.py

@@ -13,7 +13,11 @@
 
 
 class TopologicalSortPass(ir.passes.InPlacePass):
-    """Topologically sort graphs and functions in a model."""
+    """Topologically sort graphs and functions in a model.
+
+    The sort is stable, preserving the relative order of nodes that are not
+    dependent on each other. Read more at :meth:`onnx_ir.Graph.sort`.
+    """
 
     def call(self, model: ir.Model) -> ir.passes.PassResult:
         original_nodes = list(model.graph)