Use OpPattern in tracks

Author: ricardoV94

pytensor/graph/rewriting/kanren.py

@@ -74,7 +74,7 @@ def results_filter(
         self.node_filter = node_filter
         super().__init__()
 
-    def transform(self, fgraph, node):
+    def transform(self, fgraph, node, enforce_tracks: bool = True):
         if self.node_filter(node) is False:
             return False
 

pytensor/scalar/basic.py

@@ -1228,6 +1228,8 @@ def __init__(self, output_types_preference=None, name=None):
                     f"(got: {output_types_preference})"
                 )
             self.output_types_preference = output_types_preference
+        elif not hasattr(self, "output_types_preference"):
+            self.output_types_preference = None
 
     def make_node(self, *inputs):
         if self.nin >= 0:
@@ -1247,7 +1249,7 @@ def make_node(self, *inputs):
         return Apply(self, inputs, outputs)
 
     def output_types(self, types):
-        if hasattr(self, "output_types_preference"):
+        if self.output_types_preference is not None:
             variables = self.output_types_preference(*types)
             if not isinstance(variables, list | tuple) or any(
                 not isinstance(x, CType) for x in variables
@@ -2696,7 +2698,7 @@ class Sign(UnaryScalarOp):
     nfunc_spec = ("sign", 1, 1)
 
     @staticmethod
-    def output_types_preference(x):
+    def _output_types_preference(x):
         if x == bool:
             raise TypeError(x)
         return same_out_nocomplex(x)
@@ -2737,7 +2739,7 @@ def c_code_cache_version(self):
             return s
 
 
-sign = Sign(name="sign")
+sign = Sign(name="sign", output_types_preference=Sign._output_types_preference)
 
 
 class Ceil(UnaryScalarOp):

pytensor/tensor/_linalg/solve/rewriting.py

@@ -14,6 +14,7 @@
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.rewriting.basic import register_specialize
+from pytensor.tensor.rewriting.blockwise import blockwise_of
 from pytensor.tensor.rewriting.linalg import is_matrix_transpose
 from pytensor.tensor.slinalg import Solve, cho_solve, cholesky, lu_factor, lu_solve
 from pytensor.tensor.variable import TensorVariable
@@ -227,7 +228,7 @@ def _scan_split_non_sequence_decomposition_and_solve(
 
 
 @register_specialize
-@node_rewriter([Blockwise])
+@node_rewriter([blockwise_of(Solve)])
 def reuse_decomposition_multiple_solves(fgraph, node):
     return _split_decomp_and_solve_steps(
         fgraph, node, eager=False, allowed_assume_a={"gen", "tridiagonal", "pos"}

pytensor/tensor/rewriting/basic.py

@@ -29,7 +29,7 @@
 import pytensor.scalar.basic as ps
 from pytensor import compile, config
 from pytensor.compile.ops import ViewOp
-from pytensor.graph import FunctionGraph
+from pytensor.graph import FunctionGraph, Op
 from pytensor.graph.basic import Constant
 from pytensor.graph.rewriting.basic import (
     NodeProcessingGraphRewriter,
@@ -40,9 +40,9 @@
     node_rewriter,
 )
 from pytensor.graph.rewriting.db import RewriteDatabase
+from pytensor.graph.rewriting.unify import OpPattern
 from pytensor.npy_2_compat import normalize_axis_index
 from pytensor.raise_op import Assert, CheckAndRaise, assert_op
-from pytensor.scalar.basic import Second
 from pytensor.tensor.basic import (
     Alloc,
     AllocEmpty,
@@ -225,6 +225,12 @@ def register(inner_rewriter: RewriteDatabase | Rewriter):
         return node_rewriter
 
 
+def elemwise_of(scalar_op) -> OpPattern:
+    if not isinstance(scalar_op, Op | OpPattern):
+        scalar_op = OpPattern(scalar_op)
+    return OpPattern(Elemwise, scalar_op=scalar_op)
+
+
 @register_canonicalize
 @register_specialize
 @node_rewriter([TensorFromScalar])
@@ -330,9 +336,6 @@ def local_fill_sink(fgraph, node):
     f(fill(a, b), fill(c, d), e) -> fill(c, fill(a, f(b, d, e)))
     f need to be an elemwise that isn't a fill.
     """
-    if isinstance(node.op.scalar_op, Second):
-        return False
-
     models = []
     inputs = []
     for inp in node.inputs:
@@ -653,7 +656,7 @@ def local_alloc_unary(fgraph, node):
 
 @register_canonicalize
 @register_specialize
-@node_rewriter([Elemwise])
+@node_rewriter([elemwise_of(ps.Cast)])
 def local_cast_cast(fgraph, node):
     """cast(cast(x, dtype1), dtype2)
 
@@ -663,8 +666,6 @@ def local_cast_cast(fgraph, node):
           and the first cast cause an upcast.
 
     """
-    if not (isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, ps.Cast)):
-        return
     x = node.inputs[0]
     if not (
         x.owner
@@ -1031,19 +1032,13 @@ def local_useless_switch(fgraph, node):
 
 
 @register_canonicalize
-@node_rewriter([Elemwise])
+@node_rewriter([elemwise_of(ps.BinaryScalarOp | ps.Add | ps.Mul)])
 def local_merge_switch_same_cond(fgraph, node):
     """
     Merge add/sub/mul/div/minimum/maximum/... of switches sharing the same
     condition, to enable further simplification of their branches
     Example: switch(c, a, b) + switch(c, x, y) -> switch(c, a+x, b+y)
     """
-    # node must be binary elemwise or add or mul
-    if not (
-        isinstance(node.op, Elemwise)
-        and isinstance(node.op.scalar_op, ps.BinaryScalarOp | ps.Add | ps.Mul)
-    ):
-        return
     # all inputs must be switch
     if not all(
         s.owner
@@ -1174,10 +1169,9 @@ def constant_folding(fgraph, node):
 @register_infer_shape
 @register_canonicalize("fast_compile")
 @register_useless("fast_compile")
-@node_rewriter(None)
+@node_rewriter([ViewOp])
 def local_view_op(fgraph, node):
-    if isinstance(node.op, ViewOp):
-        return node.inputs
+    return node.inputs
 
 
 @register_infer_shape

pytensor/tensor/rewriting/blockwise.py

@@ -1,8 +1,9 @@
 from pytensor.compile.mode import optdb
-from pytensor.graph import Constant, node_rewriter
+from pytensor.graph import Constant, Op, node_rewriter
 from pytensor.graph.destroyhandler import inplace_candidates
 from pytensor.graph.replace import vectorize_node
 from pytensor.graph.rewriting.basic import copy_stack_trace, out2in
+from pytensor.graph.rewriting.unify import OpPattern
 from pytensor.tensor.basic import Alloc, ARange, alloc, shape_padleft
 from pytensor.tensor.blockwise import Blockwise, _squeeze_left
 from pytensor.tensor.math import Dot
@@ -20,6 +21,12 @@
 )
 
 
+def blockwise_of(core_op) -> OpPattern:
+    if not isinstance(core_op, Op | OpPattern):
+        core_op = OpPattern(core_op)
+    return OpPattern(Blockwise, core_op=core_op)
+
+
 @node_rewriter([Blockwise])
 def local_useless_blockwise(fgraph, node):
     """
@@ -71,22 +78,24 @@ def local_useless_unbatched_blockwise(fgraph, node):
 @register_canonicalize
 @register_stabilize
 @register_specialize
-@node_rewriter(tracks=[Blockwise])
+@node_rewriter(
+    tracks=[
+        blockwise_of(
+            Dot
+            | Alloc
+            | ARange
+            | Subtensor
+            | AdvancedSubtensor
+            | AdvancedIncSubtensor
+            | Reshape
+        )
+    ]
+)
 def local_eager_useless_unbatched_blockwise(fgraph, node):
-    if isinstance(
-        node.op.core_op,
-        Dot
-        | Alloc
-        | ARange
-        | Subtensor
-        | AdvancedSubtensor
-        | AdvancedIncSubtensor
-        | Reshape,
-    ):
-        # Many Dot-related rewrites (eg, all of BlasOpt) happen before specialize
-        # These other Ops can't always be trivially vectorized at runtime,
-        # since their inputs may imply non-rectangular shapes.
-        return local_useless_unbatched_blockwise.fn(fgraph, node)
+    # Many Dot-related rewrites (eg, all of BlasOpt) happen before specialize
+    # These other Ops can't always be trivially vectorized at runtime,
+    # since their inputs may imply non-rectangular shapes.
+    return local_useless_unbatched_blockwise.fn(fgraph, node)
 
 
 @register_specialize("shape_unsafe")
@@ -204,7 +213,7 @@ def local_blockwise_alloc(fgraph, node):
 
 
 @register_specialize
-@node_rewriter([Blockwise])
+@node_rewriter([blockwise_of(Reshape)])
 def local_blockwise_reshape(fgraph, node):
     """Rewrite away square Blockwise reshapes.
 
@@ -215,9 +224,6 @@ def local_blockwise_reshape(fgraph, node):
     For the square Reshape case, we must wait for all the intermediate
     operations to be lifted as Allocs
     """
-    if not isinstance(node.op.core_op, Reshape):
-        return None
-
     x, output_shape = node.inputs
     batch_ndim = node.op.batch_ndim(node)
     if all(output_shape.type.broadcastable[:batch_ndim]):

pytensor/tensor/rewriting/elemwise.py

@@ -26,6 +26,7 @@
     out2in,
 )
 from pytensor.graph.rewriting.db import SequenceDB
+from pytensor.graph.rewriting.unify import OpPattern
 from pytensor.graph.utils import InconsistencyError, MethodNotDefined
 from pytensor.scalar.math import Grad2F1Loop, _grad_2f1_loop
 from pytensor.tensor.basic import (
@@ -37,6 +38,7 @@
 from pytensor.tensor.rewriting.basic import (
     alloc_like,
     broadcasted_by,
+    elemwise_of,
     register_canonicalize,
     register_specialize,
     register_stabilize,
@@ -422,7 +424,14 @@ def local_useless_dimshuffle_makevector(fgraph, node):
 
 
 @register_canonicalize
-@node_rewriter([Elemwise])
+@node_rewriter(
+    [
+        elemwise_of(
+            OpPattern(ps.ScalarOp, output_types_preference=ps.upgrade_to_float)
+        ),
+        elemwise_of(OpPattern(ps.ScalarOp, output_types_preference=ps.upcast_out)),
+    ]
+)
 def local_upcast_elemwise_constant_inputs(fgraph, node):
     """This explicitly upcasts constant inputs to elemwise Ops, when
     those Ops do implicit upcasting anyway.
@@ -433,12 +442,6 @@ def local_upcast_elemwise_constant_inputs(fgraph, node):
     if len(node.outputs) > 1:
         return None
 
-    if getattr(node.op.scalar_op, "output_types_preference", None) not in (
-        ps.upgrade_to_float,
-        ps.upcast_out,
-    ):
-        return None
-
     # this is the kind of op that we can screw with the input
     # dtypes by upcasting explicitly
     [old_out] = node.outputs
@@ -988,13 +991,9 @@ def print_profile(stream, prof, level=0):
 
 @register_canonicalize
 @register_specialize
-@node_rewriter([Elemwise])
+@node_rewriter([elemwise_of(ps.Composite)])
 def local_useless_composite_outputs(fgraph, node):
     """Remove inputs and outputs of Composite Ops that are not used anywhere."""
-    if not (
-        isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, ps.Composite)
-    ):
-        return
     comp = node.op.scalar_op
     used_outputs_idxs = [
         i for i, o_extern in enumerate(node.outputs) if fgraph.clients[o_extern]
@@ -1104,14 +1103,10 @@ def local_careduce_fusion(fgraph, node):
     return [new_car_op(*elm_inputs)]
 
 
-@node_rewriter([Elemwise])
+@node_rewriter([elemwise_of(ps.Composite)])
 def local_inline_composite_constants(fgraph, node):
     """Inline scalar constants in Composite graphs."""
     composite_op = node.op.scalar_op
-
-    if not isinstance(composite_op, ps.Composite):
-        return None
-
     new_outer_inputs = []
     new_inner_inputs = []
     inner_replacements = {}
@@ -1287,14 +1282,9 @@ def _rebuild_partial_2f1grad_loop(node, wrt):
 
 
 @register_specialize
-@node_rewriter([Elemwise])
+@node_rewriter([elemwise_of(Grad2F1Loop)])
 def local_useless_2f1grad_loop(fgraph, node):
     # Remove unused terms from the hyp2f1 grad loop
-
-    loop_op = node.op.scalar_op
-    if not isinstance(loop_op, Grad2F1Loop):
-        return
-
     grad_related_vars = node.outputs[:-4]
     # Rewrite was already applied
     if len(grad_related_vars) // 3 != 3:
@@ -1326,18 +1316,13 @@ def local_useless_2f1grad_loop(fgraph, node):
     return replacements
 
 
-@node_rewriter([Elemwise])
+@node_rewriter([elemwise_of(Grad2F1Loop)])
 def split_2f1grad_loop(fgraph, node):
     """
     2f1grad loop has too many operands for Numpy frompyfunc code used by Elemwise nodes on python mode.
 
     This rewrite splits it across 3 different operations. It is not needed if `local_useless_2f1grad_loop` was applied
     """
-    loop_op = node.op.scalar_op
-
-    if not isinstance(loop_op, Grad2F1Loop):
-        return None
-
     grad_related_vars = node.outputs[:-4]
     # local_useless_2f1grad_loop was used, we should be safe
     if len(grad_related_vars) // 3 != 3: