Improve rewrites

Author: ricardoV94

pytensor/graph/basic.py

@@ -875,23 +875,26 @@ def walk(
     else:
         nodes_pop = nodes.pop
 
-    while nodes:
-        node: T = nodes_pop()
-
-        node_hash: int = hash_fn(node)
-
-        if node_hash not in rval_set:
-            rval_set.add(node_hash)
-
-            new_nodes: Iterable[T] | None = expand(node)
-
-            if return_children:
+    if return_children:
+        while nodes:
+            node: T = nodes_pop()
+            node_hash: int = hash_fn(node)
+            if node_hash not in rval_set:
+                new_nodes: Iterable[T] | None = expand(node)
                 yield node, new_nodes
-            else:
+                rval_set.add(node_hash)
+                if new_nodes:
+                    nodes.extend(new_nodes)
+    else:
+        while nodes:
+            node: T = nodes_pop()
+            node_hash: int = hash_fn(node)
+            if node_hash not in rval_set:
+                new_nodes: Iterable[T] | None = expand(node)
                 yield node
-
-            if new_nodes:
-                nodes.extend(new_nodes)
+                rval_set.add(node_hash)
+                if new_nodes:
+                    nodes.extend(new_nodes)
 
 
 def ancestors(

pytensor/graph/rewriting/basic.py

@@ -25,6 +25,7 @@
     applys_between,
     io_toposort,
     vars_between,
+    walk,
 )
 from pytensor.graph.features import AlreadyThere, Feature
 from pytensor.graph.fg import FunctionGraph, Output
@@ -1821,12 +1822,13 @@ class WalkingGraphRewriter(NodeProcessingGraphRewriter):
     def __init__(
         self,
         node_rewriter: NodeRewriter,
-        order: Literal["out_to_in", "in_to_out"] = "in_to_out",
+        order: Literal["out_to_in", "in_to_out", "bfs", "dfs"] = "bfs",
         ignore_newtrees: bool = False,
         failure_callback: FailureCallbackType | None = None,
     ):
-        if order not in ("out_to_in", "in_to_out"):
-            raise ValueError("order must be 'out_to_in' or 'in_to_out'")
+        valid_options = ("out_to_in", "in_to_out", "bfs", "dfs")
+        if order not in valid_options:
+            raise ValueError(f"order must be one of {valid_options}, got {order}")
         self.order = order
         super().__init__(node_rewriter, ignore_newtrees, failure_callback)
 
@@ -1836,9 +1838,19 @@ 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))
+        if self.order in {"bfs", "dfs"}:
+            q = deque(
+                walk(
+                    [out.owner for out in fgraph.outputs if out.owner is not None],
+                    expand=lambda n: (i.owner for i in n.inputs if i.owner is not None),
+                    bfs=(self.order == "bfs"),
+                )
+            )
+        else:
+            q = deque(io_toposort(fgraph.inputs, start_from))
         io_t = time.perf_counter() - t0
 
+        # Importer is ignored if self.ignore_newtrees = False
         def importer(node):
             if node is not current_node:
                 q.append(node)
@@ -1959,6 +1971,8 @@ def walking_rewriter(
 
 in2out = partial(walking_rewriter, "in_to_out")
 out2in = partial(walking_rewriter, "out_to_in")
+bfs = partial(walking_rewriter, "bfs")
+dfs = partial(walking_rewriter, "dfs")
 
 
 class ChangeTracker(Feature):

pytensor/graph/rewriting/db.py

@@ -310,7 +310,10 @@ class EquilibriumDB(RewriteDatabase):
     """
 
     def __init__(
-        self, ignore_newtrees: bool = True, tracks_on_change_inputs: bool = False
+        self,
+        ignore_newtrees: bool = True,
+        tracks_on_change_inputs: bool = False,
+        eq_rewriter_class=pytensor_rewriting.EquilibriumGraphRewriter,
     ):
         """
 
@@ -329,6 +332,7 @@ def __init__(
         self.tracks_on_change_inputs = tracks_on_change_inputs
         self.__final__: dict[str, bool] = {}
         self.__cleanup__: dict[str, bool] = {}
+        self.eq_rewriter_class = eq_rewriter_class
 
     def register(
         self,
@@ -360,7 +364,7 @@ def query(self, *tags, **kwtags):
             final_rewriters = None
         if len(cleanup_rewriters) == 0:
             cleanup_rewriters = None
-        return pytensor_rewriting.EquilibriumGraphRewriter(
+        return self.eq_rewriter_class(
             rewriters,
             max_use_ratio=config.optdb__max_use_ratio,
             ignore_newtrees=self.ignore_newtrees,

pytensor/scan/rewriting.py

@@ -30,6 +30,7 @@
 from pytensor.graph.op import compute_test_value
 from pytensor.graph.replace import clone_replace
 from pytensor.graph.rewriting.basic import (
+    EquilibriumGraphRewriter,
     GraphRewriter,
     copy_stack_trace,
     in2out,
@@ -2517,6 +2518,15 @@ def scan_push_out_dot1(fgraph, node):
     return False
 
 
+class ScanEquilibriumGraphRewriter(EquilibriumGraphRewriter):
+    """Subclass of EquilibriumGraphRewriter that aborts early if there are no Scan Ops in the graph"""
+
+    def apply(self, fgraph, start_from=None):
+        if not any(isinstance(node.op, Scan) for node in fgraph.apply_nodes):
+            return
+        super().apply(fgraph=fgraph, start_from=start_from)
+
+
 # I've added an equilibrium because later scan optimization in the sequence
 # can make it such that earlier optimizations should apply. However, in
 # general I do not expect the sequence to run more then once