Use OpPattern in tracks

Author: ricardoV94

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

@@ -26,10 +26,9 @@
 
 import numpy as np
 
-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 +39,24 @@
     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.scalar import (
+    AND,
+    EQ,
+    LE,
+    NEQ,
+    OR,
+    XOR,
+    Add,
+    BinaryScalarOp,
+    Cast,
+    Identity,
+    Mul,
+    Second,
+    Switch,
+)
 from pytensor.tensor.basic import (
     Alloc,
     AllocEmpty,
@@ -225,6 +239,12 @@ def register(inner_rewriter: RewriteDatabase | Rewriter):
         return node_rewriter
 
 
+def elemwise_of(scalar_op: type[Op] | OpPattern) -> 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])
@@ -551,15 +571,15 @@ def local_useless_elemwise(fgraph, node):
     dtype = node.outputs[0].type.dtype
     scalar_op = node.op.scalar_op
 
-    if isinstance(scalar_op, ps.EQ) and len(node.inputs) == 2:
+    if isinstance(scalar_op, EQ) and len(node.inputs) == 2:
         if node.inputs[0] is node.inputs[1]:
             # it is the same var in the graph. That will always be true
             ret = ones_like(node.inputs[0], dtype=dtype, opt=True)
 
             # Copy stack trace from input to constant output
             copy_stack_trace(node.outputs[0], ret)
             return [ret]
-    elif isinstance(scalar_op, ps.NEQ | ps.XOR) and len(node.inputs) == 2:
+    elif isinstance(scalar_op, NEQ | XOR) and len(node.inputs) == 2:
         if node.inputs[0] is node.inputs[1]:
             # it is the same var in the graph. That will always be false
             ret = zeros_like(node.inputs[0], dtype=dtype, opt=True)
@@ -568,14 +588,11 @@ def local_useless_elemwise(fgraph, node):
             copy_stack_trace(node.outputs[0], ret)
             return [ret]
 
-    elif (
-        isinstance(node.op.scalar_op, ps.Mul | ps.Add | ps.Identity)
-        and len(node.inputs) == 1
-    ):
+    elif isinstance(node.op.scalar_op, Mul | Add | Identity) and len(node.inputs) == 1:
         # No need to copy over any stack trace
         return [node.inputs[0]]
 
-    elif isinstance(node.op.scalar_op, ps.AND) and len(node.inputs) == 2:
+    elif isinstance(node.op.scalar_op, AND) and len(node.inputs) == 2:
         if (
             isinstance(node.inputs[0], TensorConstant)
             and node.inputs[1].type.broadcastable == out_bcast
@@ -602,7 +619,7 @@ def local_useless_elemwise(fgraph, node):
                     # and this rewrite would be wrong
                     return [node.inputs[0].astype(node.outputs[0].dtype)]
 
-    elif isinstance(node.op.scalar_op, ps.OR) and len(node.inputs) == 2:
+    elif isinstance(node.op.scalar_op, OR) and len(node.inputs) == 2:
         if (
             isinstance(node.inputs[0], TensorConstant)
             and node.inputs[1].type.broadcastable == out_bcast
@@ -653,7 +670,7 @@ def local_alloc_unary(fgraph, node):
 
 @register_canonicalize
 @register_specialize
-@node_rewriter([Elemwise])
+@node_rewriter([elemwise_of(Cast)])
 def local_cast_cast(fgraph, node):
     """cast(cast(x, dtype1), dtype2)
 
@@ -663,13 +680,11 @@ 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
         and isinstance(x.owner.op, Elemwise)
-        and isinstance(x.owner.op.scalar_op, ps.Cast)
+        and isinstance(x.owner.op.scalar_op, Cast)
     ):
         return
 
@@ -1009,7 +1024,7 @@ def local_useless_switch(fgraph, node):
         node.outputs[0].type.ndim == 0
         and cond_var.owner
         and isinstance(cond_var.owner.op, Elemwise)
-        and isinstance(cond_var.owner.op.scalar_op, ps.LE)
+        and isinstance(cond_var.owner.op.scalar_op, LE)
         and cond_var.owner.inputs[0].owner
         and isinstance(cond_var.owner.inputs[0].owner.op, Shape_i)
         and get_scalar_constant_value(
@@ -1031,24 +1046,18 @@ def local_useless_switch(fgraph, node):
 
 
 @register_canonicalize
-@node_rewriter([Elemwise])
+@node_rewriter([elemwise_of(BinaryScalarOp | Add | 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
         and isinstance(s.owner.op, Elemwise)
-        and isinstance(s.owner.op.scalar_op, ps.Switch)
+        and isinstance(s.owner.op.scalar_op, Switch)
         for s in node.inputs
     ):
         return
@@ -1174,10 +1183,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: type[Op] | OpPattern) -> 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]):