Faster graph traversal functions

* Avoid reversing inputs as we traverse graph
* Simplify io_toposort without ordering (and refactor into its own function)
* Removes client side-effect on previous toposort functions
* Remove duplicated logic across methods

Author: ricardoV94

pytensor/graph/basic.py

@@ -5,7 +5,6 @@
 from collections.abc import (
     Hashable,
     Iterable,
-    Reversible,
     Sequence,
 )
 from copy import copy
@@ -961,7 +960,7 @@ def clone_node_and_cache(
 
 def clone_get_equiv(
     inputs: Iterable[Variable],
-    outputs: Reversible[Variable],
+    outputs: Iterable[Variable],
     copy_inputs: bool = True,
     copy_orphans: bool = True,
     memo: dict[Union[Apply, Variable, "Op"], Union[Apply, Variable, "Op"]]
@@ -1002,7 +1001,7 @@ def clone_get_equiv(
         Keywords passed to `Apply.clone_with_new_inputs`.
 
     """
-    from pytensor.graph.traversal import io_toposort
+    from pytensor.graph.traversal import toposort
 
     if memo is None:
         memo = {}
@@ -1018,7 +1017,7 @@ def clone_get_equiv(
             memo.setdefault(input, input)
 
     # go through the inputs -> outputs graph cloning as we go
-    for apply in io_toposort(inputs, outputs):
+    for apply in toposort(outputs, blockers=inputs):
         for input in apply.inputs:
             if input not in memo:
                 if not isinstance(input, Constant) and copy_orphans:

pytensor/graph/features.py

@@ -10,7 +10,7 @@
 import pytensor
 from pytensor.configdefaults import config
 from pytensor.graph.basic import Variable
-from pytensor.graph.traversal import io_toposort
+from pytensor.graph.traversal import toposort
 from pytensor.graph.utils import InconsistencyError
 
 
@@ -340,11 +340,11 @@ def clone(self):
 
 class Bookkeeper(Feature):
     def on_attach(self, fgraph):
-        for node in io_toposort(fgraph.inputs, fgraph.outputs):
+        for node in toposort(fgraph.outputs):
             self.on_import(fgraph, node, "on_attach")
 
     def on_detach(self, fgraph):
-        for node in io_toposort(fgraph.inputs, fgraph.outputs):
+        for node in toposort(fgraph.outputs):
             self.on_prune(fgraph, node, "Bookkeeper.detach")
 
 

pytensor/graph/fg.py

@@ -19,7 +19,8 @@
 from pytensor.graph.traversal import (
     applys_between,
     graph_inputs,
-    io_toposort,
+    toposort,
+    toposort_with_orderings,
     vars_between,
 )
 from pytensor.graph.utils import MetaObject, MissingInputError, TestValueError
@@ -366,7 +367,7 @@ def import_node(
         # new nodes, so we use all variables we know of as if they were the
         # input set.  (The functions in the graph module only use the input set
         # to know where to stop going down.)
-        new_nodes = io_toposort(self.variables, apply_node.outputs)
+        new_nodes = tuple(toposort(apply_node.outputs, blockers=self.variables))
 
         if check:
             for node in new_nodes:
@@ -759,7 +760,7 @@ def toposort(self) -> list[Apply]:
             # No sorting is necessary
             return list(self.apply_nodes)
 
-        return io_toposort(self.inputs, self.outputs, self.orderings())
+        return list(toposort_with_orderings(self.outputs, orderings=self.orderings()))
 
     def orderings(self) -> dict[Apply, list[Apply]]:
         """Return a map of node to node evaluation dependencies.

pytensor/graph/replace.py

@@ -10,7 +10,10 @@
 )
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import Op
-from pytensor.graph.traversal import io_toposort, truncated_graph_inputs
+from pytensor.graph.traversal import (
+    toposort,
+    truncated_graph_inputs,
+)
 
 
 ReplaceTypes = Iterable[tuple[Variable, Variable]] | dict[Variable, Variable]
@@ -295,7 +298,7 @@ def vectorize_graph(
     new_inputs = [replace.get(inp, inp) for inp in inputs]
 
     vect_vars = dict(zip(inputs, new_inputs, strict=True))
-    for node in io_toposort(inputs, seq_outputs):
+    for node in toposort(seq_outputs, blockers=inputs):
         vect_inputs = [vect_vars.get(inp, inp) for inp in node.inputs]
         vect_node = vectorize_node(node, *vect_inputs)
         for output, vect_output in zip(node.outputs, vect_node.outputs, strict=True):

pytensor/graph/rewriting/basic.py

@@ -27,7 +27,7 @@
 from pytensor.graph.fg import FunctionGraph, Output
 from pytensor.graph.op import Op
 from pytensor.graph.rewriting.unify import Var, convert_strs_to_vars
-from pytensor.graph.traversal import applys_between, io_toposort, vars_between
+from pytensor.graph.traversal import applys_between, toposort, vars_between
 from pytensor.graph.utils import AssocList, InconsistencyError
 from pytensor.misc.ordered_set import OrderedSet
 from pytensor.utils import flatten
@@ -1834,7 +1834,7 @@ def apply(self, fgraph, start_from=None):
         callback_before = fgraph.execute_callbacks_time
         nb_nodes_start = len(fgraph.apply_nodes)
         t0 = time.perf_counter()
-        q = deque(io_toposort(fgraph.inputs, start_from))
+        q = deque(toposort(start_from))
         io_t = time.perf_counter() - t0
 
         def importer(node):
@@ -2164,7 +2164,7 @@ def apply_cleanup(profs_dict):
             changed |= apply_cleanup(iter_cleanup_sub_profs)
 
             topo_t0 = time.perf_counter()
-            q = deque(io_toposort(fgraph.inputs, start_from))
+            q = deque(toposort(start_from))
             io_toposort_timing.append(time.perf_counter() - topo_t0)
 
             nb_nodes.append(len(q))