Replace uses of in2out and out2in by a depth-first search 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 applys_between, toposort, vars_between
+from pytensor.graph.traversal import (
+    apply_ancestors,
+    applys_between,
+    toposort,
+    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", "dfs"] = "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", "dfs")
+        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,11 @@ 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(toposort(start_from))
+        q = deque(
+            apply_ancestors(start_from)
+            if (self.order == "dfs")
+            else toposort(start_from)
+        )
         io_t = time.perf_counter() - t0
 
         def importer(node):
@@ -1957,6 +1967,7 @@ def walking_rewriter(
 
 in2out = partial(walking_rewriter, "in_to_out")
 out2in = partial(walking_rewriter, "out_to_in")
+dfs_rewriter = partial(walking_rewriter, "dfs")
 
 
 class ChangeTracker(Feature):

pytensor/scan/rewriting.py

@@ -28,7 +28,7 @@
 from pytensor.graph.rewriting.basic import (
     GraphRewriter,
     copy_stack_trace,
-    in2out,
+    dfs_rewriter,
     node_rewriter,
 )
 from pytensor.graph.rewriting.db import EquilibriumDB, SequenceDB
@@ -2548,15 +2548,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),
+    dfs_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),
+    dfs_rewriter(scan_save_mem_no_prealloc, ignore_newtrees=True),
     "numba",
     "jax",
     "pytorch",
@@ -2577,7 +2577,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),
+    dfs_rewriter(remove_constants_and_unused_inputs_scan, ignore_newtrees=True),
     "remove_constants_and_unused_inputs_scan",
     "fast_run",
     "scan",
@@ -2586,7 +2586,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),
+    dfs_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",
@@ -2596,7 +2596,7 @@ def scan_push_out_dot1(fgraph, node):
 
 scan_seqopt1.register(
     "scan_push_out_seq",
-    in2out(scan_push_out_seq, ignore_newtrees=True),
+    dfs_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",
@@ -2607,7 +2607,7 @@ def scan_push_out_dot1(fgraph, node):
 
 scan_seqopt1.register(
     "scan_push_out_dot1",
-    in2out(scan_push_out_dot1, ignore_newtrees=True),
+    dfs_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",
@@ -2620,7 +2620,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),
+    dfs_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",
@@ -2631,22 +2631,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),
+    dfs_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),
+    dfs_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),
+    dfs_rewriter(remove_constants_and_unused_inputs_scan, ignore_newtrees=True),
     "remove_constants_and_unused_inputs_scan",
     "fast_run",
     "scan",
@@ -2661,23 +2661,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),
+    dfs_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),
+    dfs_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),
+    dfs_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 copy_stack_trace, dfs_rewriter, 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),
+    dfs_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),
+    dfs_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),
+    dfs_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 (
+    copy_stack_trace,
+    dfs_rewriter,
+    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),
+    dfs_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 dfs_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,16 @@ 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 = dfs_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 dfs_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),
+    dfs_rewriter(introduce_explicit_core_shape_rv),
     "numba",
     position=100,
 )

pytensor/tensor/rewriting/basic.py

@@ -36,6 +36,7 @@
     NodeRewriter,
     Rewriter,
     copy_stack_trace,
+    dfs_rewriter,
     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),
+    dfs_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

@@ -77,7 +77,7 @@
     EquilibriumGraphRewriter,
     GraphRewriter,
     copy_stack_trace,
-    in2out,
+    dfs_rewriter,
     node_rewriter,
 )
 from pytensor.graph.rewriting.db import SequenceDB
@@ -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),
+    dfs_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 = dfs_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),
+    dfs_rewriter(local_dot22_to_dot22scalar),
     "fast_run",
     position=12,
 )