Faster graph traversal functions

* Uses simpler new apply_toposort
* Removes client side-effect on previous toposort functions

Author: ricardoV94

pytensor/graph/basic.py

@@ -1000,7 +1000,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 variable_bounded_apply_toposort
 
     if memo is None:
         memo = {}
@@ -1016,7 +1016,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 variable_bounded_apply_toposort(outputs, blockers=inputs):
         for input in apply.inputs:
             if input not in memo:
                 if not isinstance(input, Constant) and copy_orphans:

pytensor/graph/fg.py

@@ -17,13 +17,13 @@
 from pytensor.graph.features import AlreadyThere, Feature, ReplaceValidate
 from pytensor.graph.op import Op
 from pytensor.graph.traversal import (
+    apply_toposort,
     applys_between,
     graph_inputs,
-    io_toposort,
+    variable_bounded_apply_toposort,
     vars_between,
 )
 from pytensor.graph.utils import MetaObject, MissingInputError, TestValueError
-from pytensor.misc.ordered_set import OrderedSet
 
 
 ClientType = tuple[Apply, int]
@@ -132,7 +132,6 @@ def __init__(
             features = []
 
         self._features: list[Feature] = []
-
         # All apply nodes in the subgraph defined by inputs and
         # outputs are cached in this field
         self.apply_nodes: set[Apply] = set()
@@ -355,21 +354,16 @@ def import_node(
         apply_node : Apply
             The node to be imported.
         check : bool
-            Check that the inputs for the imported nodes are also present in
-            the `FunctionGraph`.
+            Check that the inputs for the imported nodes are also present in the `FunctionGraph`.
         reason : str
             The name of the optimization or operation in progress.
         import_missing : bool
             Add missing inputs instead of raising an exception.
         """
         # We import the nodes in topological order. We only are interested in
-        # 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)
-
-        if check:
-            for node in new_nodes:
+        # new nodes, so we use all nodes we know of as inputs to interrupt the toposort
+        for node in apply_toposort([apply_node], blockers=self.apply_nodes):
+            if check:
                 for var in node.inputs:
                     if (
                         var.owner is None
@@ -389,8 +383,6 @@ def import_node(
                             )
                             raise MissingInputError(error_msg, variable=var)
 
-        for node in new_nodes:
-            assert node not in self.apply_nodes
             self.apply_nodes.add(node)
             if not hasattr(node.tag, "imported_by"):
                 node.tag.imported_by = []
@@ -755,11 +747,13 @@ def toposort(self) -> list[Apply]:
           :meth:`FunctionGraph.orderings`.
 
         """
-        if len(self.apply_nodes) < 2:
-            # No sorting is necessary
-            return list(self.apply_nodes)
-
-        return io_toposort(self.inputs, self.outputs, self.orderings())
+        if orderings := self.orderings():
+            return list(
+                variable_bounded_apply_toposort(self.outputs, self.inputs, orderings)
+            )
+        else:
+            # Faster implementation when no orderings are needed
+            return list(apply_toposort(o.owner for o in self.outputs))
 
     def orderings(self) -> dict[Apply, list[Apply]]:
         """Return a map of node to node evaluation dependencies.
@@ -778,29 +772,16 @@ def orderings(self) -> dict[Apply, list[Apply]]:
         take care of computing the dependencies by itself.
 
         """
-        assert isinstance(self._features, list)
         all_orderings: list[dict] = []
 
         for feature in self._features:
             if hasattr(feature, "orderings"):
-                orderings = feature.orderings(self)
-                if not isinstance(orderings, dict):
-                    raise TypeError(
-                        "Non-deterministic return value from "
-                        + str(feature.orderings)
-                        + ". Nondeterministic object is "
-                        + str(orderings)
-                    )
-                if len(orderings) > 0:
+                if orderings := feature.orderings(self):
                     all_orderings.append(orderings)
-                    for node, prereqs in orderings.items():
-                        if not isinstance(prereqs, list | OrderedSet):
-                            raise TypeError(
-                                "prereqs must be a type with a "
-                                "deterministic iteration order, or toposort "
-                                " will be non-deterministic."
-                            )
-        if len(all_orderings) == 1:
+
+        if not all_orderings:
+            return {}
+        elif len(all_orderings) == 1:
             # If there is only 1 ordering, we reuse it directly.
             return all_orderings[0].copy()
         else:

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 (
+    truncated_graph_inputs,
+    variable_bounded_apply_toposort,
+)
 
 
 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 variable_bounded_apply_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 apply_toposort, applys_between, 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(apply_toposort(o.owner for o in start_from))
         io_t = time.perf_counter() - t0
 
         def importer(node):
@@ -2081,11 +2081,6 @@ def add_requirements(self, fgraph):
     def apply(self, fgraph, start_from=None):
         change_tracker = ChangeTracker()
         fgraph.attach_feature(change_tracker)
-        if start_from is None:
-            start_from = fgraph.outputs
-        else:
-            for node in start_from:
-                assert node in fgraph.outputs
 
         changed = True
         max_use_abort = False
@@ -2164,7 +2159,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(apply_toposort(o.owner for o in fgraph.outputs))
             io_toposort_timing.append(time.perf_counter() - topo_t0)
 
             nb_nodes.append(len(q))