Replace uses of in2out and out2in by bfs_rewriter

Author: ricardoV94

pytensor/graph/rewriting/basic.py

@@ -27,7 +27,12 @@
 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 apply_toposort, applys_between, vars_between
+from pytensor.graph.traversal import (
+    apply_ancestors,
+    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
@@ -1819,12 +1824,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"] = "in_to_out",
         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_orders = ("out_to_in", "in_to_out", "bfs")
+        if order not in valid_orders:
+            raise ValueError(f"order must be one of {valid_orders}, got {order}")
         self.order = order
         super().__init__(node_rewriter, ignore_newtrees, failure_callback)
 
@@ -1834,7 +1840,10 @@ 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(apply_toposort(output_nodes=(o.owner for o in start_from)))
+        if self.order == "bfs":
+            q = deque(apply_ancestors(o.owner for o in start_from))
+        else:
+            q = deque(apply_toposort(o.owner for o in start_from))
         io_t = time.perf_counter() - t0
 
         def importer(node):
@@ -1957,6 +1966,7 @@ def walking_rewriter(
 
 in2out = partial(walking_rewriter, "in_to_out")
 out2in = partial(walking_rewriter, "out_to_in")
+bfs_rewriter = partial(walking_rewriter, "bfs")
 
 
 class ChangeTracker(Feature):

pytensor/scan/rewriting.py

@@ -27,8 +27,8 @@
 from pytensor.graph.replace import clone_replace
 from pytensor.graph.rewriting.basic import (
     GraphRewriter,
+    bfs_rewriter,
     copy_stack_trace,
-    in2out,
     node_rewriter,
 )
 from pytensor.graph.rewriting.db import EquilibriumDB, SequenceDB
@@ -2534,15 +2534,15 @@ def scan_push_out_dot1(fgraph, node):
 # ScanSaveMem should execute only once per node.
 optdb.register(
     "scan_save_mem_prealloc",
-    in2out(scan_save_mem_prealloc, ignore_newtrees=True),
+    bfs_rewriter(scan_save_mem_prealloc, ignore_newtrees=True),
     "fast_run",
     "scan",
     "scan_save_mem",
     position=1.61,
 )
 optdb.register(
     "scan_save_mem_no_prealloc",
-    in2out(scan_save_mem_no_prealloc, ignore_newtrees=True),
+    bfs_rewriter(scan_save_mem_no_prealloc, ignore_newtrees=True),
     "numba",
     "jax",
     "pytorch",
@@ -2563,7 +2563,7 @@ def scan_push_out_dot1(fgraph, node):
 
 scan_seqopt1.register(
     "scan_remove_constants_and_unused_inputs0",
-    in2out(remove_constants_and_unused_inputs_scan, ignore_newtrees=True),
+    bfs_rewriter(remove_constants_and_unused_inputs_scan, ignore_newtrees=True),
     "remove_constants_and_unused_inputs_scan",
     "fast_run",
     "scan",
@@ -2572,7 +2572,7 @@ def scan_push_out_dot1(fgraph, node):
 
 scan_seqopt1.register(
     "scan_push_out_non_seq",
-    in2out(scan_push_out_non_seq, ignore_newtrees=True),
+    bfs_rewriter(scan_push_out_non_seq, ignore_newtrees=True),
     "scan_pushout_nonseqs_ops",  # For backcompat: so it can be tagged with old name
     "fast_run",
     "scan",
@@ -2582,7 +2582,7 @@ def scan_push_out_dot1(fgraph, node):
 
 scan_seqopt1.register(
     "scan_push_out_seq",
-    in2out(scan_push_out_seq, ignore_newtrees=True),
+    bfs_rewriter(scan_push_out_seq, ignore_newtrees=True),
     "scan_pushout_seqs_ops",  # For backcompat: so it can be tagged with old name
     "fast_run",
     "scan",
@@ -2593,7 +2593,7 @@ def scan_push_out_dot1(fgraph, node):
 
 scan_seqopt1.register(
     "scan_push_out_dot1",
-    in2out(scan_push_out_dot1, ignore_newtrees=True),
+    bfs_rewriter(scan_push_out_dot1, ignore_newtrees=True),
     "scan_pushout_dot1",  # For backcompat: so it can be tagged with old name
     "fast_run",
     "more_mem",
@@ -2606,7 +2606,7 @@ def scan_push_out_dot1(fgraph, node):
 scan_seqopt1.register(
     "scan_push_out_add",
     # TODO: Perhaps this should be an `EquilibriumGraphRewriter`?
-    in2out(scan_push_out_add, ignore_newtrees=False),
+    bfs_rewriter(scan_push_out_add, ignore_newtrees=False),
     "scan_pushout_add",  # For backcompat: so it can be tagged with old name
     "fast_run",
     "more_mem",
@@ -2617,22 +2617,22 @@ def scan_push_out_dot1(fgraph, node):
 
 scan_eqopt2.register(
     "while_scan_merge_subtensor_last_element",
-    in2out(while_scan_merge_subtensor_last_element, ignore_newtrees=True),
+    bfs_rewriter(while_scan_merge_subtensor_last_element, ignore_newtrees=True),
     "fast_run",
     "scan",
 )
 
 scan_eqopt2.register(
     "constant_folding_for_scan2",
-    in2out(constant_folding, ignore_newtrees=True),
+    bfs_rewriter(constant_folding, ignore_newtrees=True),
     "fast_run",
     "scan",
 )
 
 
 scan_eqopt2.register(
     "scan_remove_constants_and_unused_inputs1",
-    in2out(remove_constants_and_unused_inputs_scan, ignore_newtrees=True),
+    bfs_rewriter(remove_constants_and_unused_inputs_scan, ignore_newtrees=True),
     "remove_constants_and_unused_inputs_scan",
     "fast_run",
     "scan",
@@ -2647,23 +2647,23 @@ def scan_push_out_dot1(fgraph, node):
 # After Merge optimization
 scan_eqopt2.register(
     "scan_remove_constants_and_unused_inputs2",
-    in2out(remove_constants_and_unused_inputs_scan, ignore_newtrees=True),
+    bfs_rewriter(remove_constants_and_unused_inputs_scan, ignore_newtrees=True),
     "remove_constants_and_unused_inputs_scan",
     "fast_run",
     "scan",
 )
 
 scan_eqopt2.register(
     "scan_merge_inouts",
-    in2out(scan_merge_inouts, ignore_newtrees=True),
+    bfs_rewriter(scan_merge_inouts, ignore_newtrees=True),
     "fast_run",
     "scan",
 )
 
 # After everything else
 scan_eqopt2.register(
     "scan_remove_constants_and_unused_inputs3",
-    in2out(remove_constants_and_unused_inputs_scan, ignore_newtrees=True),
+    bfs_rewriter(remove_constants_and_unused_inputs_scan, ignore_newtrees=True),
     "remove_constants_and_unused_inputs_scan",
     "fast_run",
     "scan",

pytensor/tensor/_linalg/solve/rewriting.py

@@ -3,7 +3,7 @@
 
 from pytensor.compile import optdb
 from pytensor.graph import Constant, graph_inputs
-from pytensor.graph.rewriting.basic import copy_stack_trace, in2out, node_rewriter
+from pytensor.graph.rewriting.basic import bfs_rewriter, copy_stack_trace, node_rewriter
 from pytensor.scan.op import Scan
 from pytensor.scan.rewriting import scan_seqopt1
 from pytensor.tensor._linalg.solve.tridiagonal import (
@@ -243,7 +243,7 @@ def scan_split_non_sequence_decomposition_and_solve(fgraph, node):
 
 scan_seqopt1.register(
     scan_split_non_sequence_decomposition_and_solve.__name__,
-    in2out(scan_split_non_sequence_decomposition_and_solve, ignore_newtrees=True),
+    bfs_rewriter(scan_split_non_sequence_decomposition_and_solve, ignore_newtrees=True),
     "fast_run",
     "scan",
     "scan_pushout",
@@ -260,7 +260,7 @@ def reuse_decomposition_multiple_solves_jax(fgraph, node):
 
 optdb["specialize"].register(
     reuse_decomposition_multiple_solves_jax.__name__,
-    in2out(reuse_decomposition_multiple_solves_jax, ignore_newtrees=True),
+    bfs_rewriter(reuse_decomposition_multiple_solves_jax, ignore_newtrees=True),
     "jax",
     use_db_name_as_tag=False,
 )
@@ -275,7 +275,9 @@ def scan_split_non_sequence_decomposition_and_solve_jax(fgraph, node):
 
 scan_seqopt1.register(
     scan_split_non_sequence_decomposition_and_solve_jax.__name__,
-    in2out(scan_split_non_sequence_decomposition_and_solve_jax, ignore_newtrees=True),
+    bfs_rewriter(
+        scan_split_non_sequence_decomposition_and_solve_jax, ignore_newtrees=True
+    ),
     "jax",
     use_db_name_as_tag=False,
     position=2,

pytensor/tensor/random/rewriting/basic.py

@@ -4,7 +4,11 @@
 from pytensor.configdefaults import config
 from pytensor.graph import ancestors
 from pytensor.graph.op import compute_test_value
-from pytensor.graph.rewriting.basic import copy_stack_trace, in2out, node_rewriter
+from pytensor.graph.rewriting.basic import (
+    bfs_rewriter,
+    copy_stack_trace,
+    node_rewriter,
+)
 from pytensor.tensor import NoneConst, TensorVariable
 from pytensor.tensor.basic import constant
 from pytensor.tensor.elemwise import DimShuffle
@@ -57,7 +61,7 @@ def random_make_inplace(fgraph, node):
 
 optdb.register(
     "random_make_inplace",
-    in2out(random_make_inplace, ignore_newtrees=True),
+    bfs_rewriter(random_make_inplace, ignore_newtrees=True),
     "fast_run",
     "inplace",
     position=50.9,

pytensor/tensor/random/rewriting/jax.py

@@ -2,8 +2,7 @@
 
 from pytensor.compile import optdb
 from pytensor.graph import Constant
-from pytensor.graph.rewriting.basic import in2out, node_rewriter
-from pytensor.graph.rewriting.db import SequenceDB
+from pytensor.graph.rewriting.basic import bfs_rewriter, in2out, node_rewriter
 from pytensor.tensor import abs as abs_t
 from pytensor.tensor import broadcast_arrays, exp, floor, log, log1p, reciprocal, sqrt
 from pytensor.tensor.basic import (
@@ -179,51 +178,18 @@ def materialize_implicit_arange_choice_without_replacement(fgraph, node):
     return new_op.make_node(rng, size, a_vector_param, *other_params).outputs
 
 
-random_vars_opt = SequenceDB()
-random_vars_opt.register(
-    "lognormal_from_normal",
-    in2out(lognormal_from_normal),
-    "jax",
-)
-random_vars_opt.register(
-    "halfnormal_from_normal",
-    in2out(halfnormal_from_normal),
-    "jax",
-)
-random_vars_opt.register(
-    "geometric_from_uniform",
-    in2out(geometric_from_uniform),
-    "jax",
-)
-random_vars_opt.register(
-    "negative_binomial_from_gamma_poisson",
-    in2out(negative_binomial_from_gamma_poisson),
-    "jax",
-)
-random_vars_opt.register(
-    "inverse_gamma_from_gamma",
-    in2out(inverse_gamma_from_gamma),
-    "jax",
-)
-random_vars_opt.register(
-    "generalized_gamma_from_gamma",
-    in2out(generalized_gamma_from_gamma),
-    "jax",
-)
-random_vars_opt.register(
-    "wald_from_normal_uniform",
-    in2out(wald_from_normal_uniform),
-    "jax",
-)
-random_vars_opt.register(
-    "beta_binomial_from_beta_binomial",
-    in2out(beta_binomial_from_beta_binomial),
-    "jax",
-)
-random_vars_opt.register(
-    "materialize_implicit_arange_choice_without_replacement",
-    in2out(materialize_implicit_arange_choice_without_replacement),
-    "jax",
+random_vars_opt = bfs_rewriter(
+    [
+        lognormal_from_normal,
+        halfnormal_from_normal,
+        geometric_from_uniform,
+        negative_binomial_from_gamma_poisson,
+        inverse_gamma_from_gamma,
+        generalized_gamma_from_gamma,
+        wald_from_normal_uniform,
+        beta_binomial_from_beta_binomial,
+        materialize_implicit_arange_choice_without_replacement,
+    ]
 )
 optdb.register("jax_random_vars_rewrites", random_vars_opt, "jax", position=110)
 

pytensor/tensor/random/rewriting/numba.py

@@ -1,6 +1,6 @@
 from pytensor.compile import optdb
 from pytensor.graph import node_rewriter
-from pytensor.graph.rewriting.basic import out2in
+from pytensor.graph.rewriting.basic import bfs_rewriter
 from pytensor.tensor import as_tensor, constant
 from pytensor.tensor.random.op import RandomVariable, RandomVariableWithCoreShape
 from pytensor.tensor.rewriting.shape import ShapeFeature
@@ -82,7 +82,7 @@ def introduce_explicit_core_shape_rv(fgraph, node):
 
 optdb.register(
     introduce_explicit_core_shape_rv.__name__,
-    out2in(introduce_explicit_core_shape_rv),
+    bfs_rewriter(introduce_explicit_core_shape_rv),
     "numba",
     position=100,
 )

pytensor/tensor/rewriting/basic.py

@@ -35,6 +35,7 @@
     NodeProcessingGraphRewriter,
     NodeRewriter,
     Rewriter,
+    bfs_rewriter,
     copy_stack_trace,
     in2out,
     node_rewriter,
@@ -518,7 +519,7 @@ def local_alloc_empty_to_zeros(fgraph, node):
 
 compile.optdb.register(
     "local_alloc_empty_to_zeros",
-    in2out(local_alloc_empty_to_zeros),
+    bfs_rewriter(local_alloc_empty_to_zeros),
     # After move to gpu and merge2, before inplace.
     "alloc_empty_to_zeros",
     position=49.3,

pytensor/tensor/rewriting/blas.py

@@ -76,8 +76,8 @@
 from pytensor.graph.rewriting.basic import (
     EquilibriumGraphRewriter,
     GraphRewriter,
+    bfs_rewriter,
     copy_stack_trace,
-    in2out,
     node_rewriter,
 )
 from pytensor.graph.rewriting.db import SequenceDB
@@ -460,7 +460,8 @@ def apply(self, fgraph):
             callbacks_before = fgraph.execute_callbacks_times.copy()
             callback_before = fgraph.execute_callbacks_time
 
-        nodelist = apply_toposort(o.owner for o in fgraph.outputs)
+        nodelist = list(apply_toposort(o.owner for o in fgraph.outputs))
+        nodelist.reverse()
 
         def on_import(new_node):
             if new_node is not node:
@@ -475,7 +476,6 @@ def on_import(new_node):
             t0 = time.perf_counter()
             time_toposort += time.perf_counter() - t0
             did_something = False
-            nodelist.reverse()
             for node in nodelist:
                 if not (
                     isinstance(node.op, Elemwise)
@@ -721,7 +721,7 @@ def local_dot22_to_ger_or_gemv(fgraph, node):
 # fast_compile is needed to have GpuDot22 created.
 blas_optdb.register(
     "local_dot_to_dot22",
-    in2out(local_dot_to_dot22),
+    bfs_rewriter(local_dot_to_dot22),
     "fast_run",
     "fast_compile",
     position=0,
@@ -744,7 +744,7 @@ def local_dot22_to_ger_or_gemv(fgraph, node):
 )
 
 
-blas_opt_inplace = in2out(
+blas_opt_inplace = bfs_rewriter(
     local_inplace_gemm, local_inplace_gemv, local_inplace_ger, name="blas_opt_inplace"
 )
 optdb.register(
@@ -883,7 +883,7 @@ def local_dot22_to_dot22scalar(fgraph, node):
 # dot22scalar and gemm give more speed up then dot22scalar
 blas_optdb.register(
     "local_dot22_to_dot22scalar",
-    in2out(local_dot22_to_dot22scalar),
+    bfs_rewriter(local_dot22_to_dot22scalar),
     "fast_run",
     position=12,
 )