Allow unifying with Op classes

Author: ricardoV94

pytensor/graph/rewriting/basic.py

@@ -29,7 +29,7 @@
 from pytensor.graph.features import AlreadyThere, Feature
 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.rewriting.unify import OpInstance, Var, convert_strs_to_vars
 from pytensor.graph.utils import AssocList, InconsistencyError
 from pytensor.misc.ordered_set import OrderedSet
 from pytensor.utils import flatten
@@ -1320,6 +1320,7 @@ class PatternNodeRewriter(NodeRewriter):
     The input and output patterns have the following syntax:
 
         input_pattern ::= (op, <sub_pattern1>, <sub_pattern2>, ...)
+        input_pattern ::= (OpInstance(type(op), {<param>: <value>, ...}), <sub_pattern1>, <sub_pattern2>, ...)
         input_pattern ::= dict(pattern = <input_pattern>,
                                constraint = <constraint>)
         sub_pattern ::= input_pattern
@@ -1333,6 +1334,7 @@ class PatternNodeRewriter(NodeRewriter):
         output_pattern ::= string
         output_pattern ::= int
         output_pattern ::= float
+        output_pattern ::= callable
 
     Each string in the input pattern is a variable that will be set to
     whatever expression is found in its place. If the same string is
@@ -1358,20 +1360,73 @@ class PatternNodeRewriter(NodeRewriter):
     Examples
     --------
 
-        PatternNodeRewriter((add, 'x', 'y'), (add, 'y', 'x'))
-        PatternNodeRewriter((multiply, 'x', 'x'), (square, 'x'))
-        PatternNodeRewriter((subtract, (add, 'x', 'y'), 'y'), 'x')
-        PatternNodeRewriter((power, 'x', Constant(double, 2.0)), (square, 'x'))
-        PatternNodeRewriter((boggle, {'pattern': 'x',
-                            'constraint': lambda expr: expr.type == scrabble}),
-                   (scrabble, 'x'))
+    .. code-block:: python
 
+        from pytensor.graph.rewriting.basic import PatternNodeRewriter
+        from pytensor.tensor import add, mul, sub, pow, square
+
+        PatternNodeRewriter((add, "x", "y"), (add, "y", "x"))
+        PatternNodeRewriter((mul, "x", "x"), (square, "x"))
+        PatternNodeRewriter((sub, (add, "x", "y"), "y"), "x")
+        PatternNodeRewriter((pow, "x", 2.0), (square, "x"))
+        PatternNodeRewriter(
+            (mul, {"pattern": "x", "constraint": lambda expr: expr.ndim == 0}, "y"),
+            (mul, "y", "x"),
+        )
+
+    You can use OpInstance to match a subtype of an Op, with some parameter constraints
+    You can also specify a callable as the output pattern, which will be called with (fgraph, node, subs_dict) as arguments.
+
+
+    .. code-block:: python
+
+        from pytensor.graph.rewriting.basic import PatternNodeRewriter
+        from pytensor.graph.rewriting.unify import OpInstance
+        from pytensor.tensor.basic import Join
+        from pytensor.tensor.elemwise import CAReduce, Elemwise
+
+
+        def output_fn(fgraph, node, s):
+            reduce_op = node.op
+            reduced_a = reduce_op(s["a"])
+            reduced_b = reduce_op(s["b"])
+            return Elemwise(s["scalar_op"])(reduced_a, reduced_b)
+
+
+        PatternNodeRewriter(
+            (
+                OpInstance(CAReduce, scalar_op="scalar_op", axis=None),
+                (Join(), "join_axis", "a", "b"),
+            ),
+            output_fn,
+        )
+
+
+    If you want to test a string parameter, you must use LiteralString to avoid it being interpreted as a unification variable.
+
+    .. code-block:: python
+
+
+        from pytensor.graph.rewriting.basic import PatternNodeRewriter
+        from pytensor.graph.rewriting.unify import OpInstance, LiteralString
+        from pytensor.tensor.blockwise import Blockwise
+        from pytensor.tensor.slinalg import Solve
+
+        PatternNodeRewriter(
+            (
+                OpInstance(
+                    Blockwise, core_op=OpInstance(Solve, assume_a=LiteralString("gen"))
+                ),
+                "A",
+                "b",
+            )
+        )
     """
 
     def __init__(
         self,
-        in_pattern,
-        out_pattern,
+        in_pattern: tuple,
+        out_pattern: tuple | Callable,
         allow_multiple_clients: bool = False,
         name: str | None = None,
         tracks=(),
@@ -1386,7 +1441,7 @@ def __init__(
         in_pattern
             The input pattern that we want to replace.
         out_pattern
-            The replacement pattern.
+            The replacement pattern. Or a callable that takes (fgraph, node, subs_dict) as inputs
         allow_multiple_clients
             If ``False``, the pattern matching will fail if one of the subpatterns has
             more than one client.
@@ -1415,26 +1470,35 @@ def __init__(
         self.out_pattern = convert_strs_to_vars(out_pattern, var_map=var_map)
         self.values_eq_approx = values_eq_approx
         self.allow_cast = allow_cast
-        if isinstance(in_pattern, list | tuple):
-            self.op = self.in_pattern[0]
-        elif isinstance(in_pattern, dict):
-            self.op = self.in_pattern["pattern"][0]
-        else:
-            raise TypeError(
-                "The pattern to search for must start with a specific Op instance."
-            )
         self.allow_multiple_clients = allow_multiple_clients
         if name:
             self.__name__ = name
-        self._tracks = tracks
         self.get_nodes = get_nodes
         if tracks != ():
-            assert get_nodes
+            if not get_nodes:
+                raise ValueError("Custom `tracks` requires `get_nodes` to be provided.")
+            self._tracks = tracks
+        else:
+            if isinstance(in_pattern, list | tuple):
+                op = self.in_pattern[0]
+            elif isinstance(in_pattern, dict):
+                op = self.in_pattern["pattern"][0]
+            else:
+                raise TypeError(
+                    "The pattern to search for must start with a specific Op instance."
+                )
+            if isinstance(op, Op):
+                self._tracks = [op]
+            elif isinstance(op, OpInstance):
+                self._tracks = [op.op_type]
+            else:
+                raise ValueError(
+                    f"The pattern to search for must start with a specific Op instance or an OpInstance class. "
+                    f"Got {op}, with type {type(op)}."
+                )
 
     def tracks(self):
-        if self._tracks != ():
-            return self._tracks
-        return [self.op]
+        return self._tracks
 
     def transform(self, fgraph, node, get_nodes=True):
         """Check if the graph from node corresponds to ``in_pattern``.
@@ -1455,28 +1519,39 @@ def transform(self, fgraph, node, get_nodes=True):
             # PatternNodeRewriter doesn't support replacing multi-output nodes
             return False
 
-        s = unify(self.in_pattern, node.out)
+        s = unify(self.in_pattern, node.out, {})
 
         if s is False:
             return False
 
-        ret = reify(self.out_pattern, s)
-
-        if isinstance(ret, ExpressionTuple):
-            ret = ret.evaled_obj
-
-        if self.values_eq_approx:
-            ret.tag.values_eq_approx = self.values_eq_approx
-
         if not self.allow_multiple_clients:
-            input_vars = list(s.values())
+            input_vars = set(s.values())
+            clients = fgraph.clients
             if any(
-                len(fgraph.clients[v]) > 1
+                len(clients[v]) > 1
                 for v in vars_between(input_vars, node.inputs)
                 if v not in input_vars
             ):
                 return False
 
+        if callable(self.out_pattern):
+            # token is the variable name used in the original pattern
+            ret = self.out_pattern(fgraph, node, {k.token: v for k, v in s.items()})
+            if ret is None or ret is False:
+                # The output function is still allowed to reject the rewrite
+                return False
+            if not isinstance(ret, Variable):
+                raise ValueError(
+                    f"The output of the PatternNodeRewriter callable must be a variable got {ret} of type {type(ret)}."
+                )
+        else:
+            ret = reify(self.out_pattern, s)
+            if isinstance(ret, ExpressionTuple):
+                ret = ret.evaled_obj
+
+        if self.values_eq_approx:
+            ret.tag.values_eq_approx = self.values_eq_approx
+
         [old_out] = node.outputs
         if not old_out.type.is_super(ret.type):
             from pytensor.tensor.type import TensorType

pytensor/graph/rewriting/unify.py

@@ -10,8 +10,10 @@
 
 """
 
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
 from numbers import Number
+from typing import Any
 
 import numpy as np
 from cons.core import ConsError, _car, _cdr
@@ -254,6 +256,103 @@ def _unify_ConstrainedVar_object(u, v, s):
 _unify.add((object, ConstrainedVar, Mapping), _unify_ConstrainedVar_object)
 
 
+@dataclass(frozen=True)
+class LiteralString:
+    value: str
+
+
+class OpInstance:
+    """Class that can be unified with Op instances of a given type and parameters.
+
+    An op instance is unified as long as the parameters specified in the OpInstance can be unified as well.
+    Parameters that are not specified in the OpInstance are ignored during unification.
+
+    This is needed because some Ops can be complex to parametrize fully,
+    and not all parameters are relevant for a given pattern.
+
+    Examples
+    --------
+
+    .. testcode::
+
+        from unification import var, unify
+        from etuples import etuple
+
+        import pytensor.tensor as pt
+        from pytensor.graph.rewriting.unify import OpInstance
+        from pytensor.tensor.blockwise import Blockwise
+        from pytensor.tensor.slinalg import Solve
+
+        A = var("A")
+        b = var("b")
+        pattern = etuple(
+            OpInstance(Blockwise, core_op=OpInstance(Solve, assume_a="gen")), A, b
+        )
+
+        A_pt = pt.tensor3("A")
+        b_pt = pt.tensor3("b")
+        out1 = pt.linalg.solve(A_pt, b_pt)
+        out2 = pt.linalg.solve(A_pt, b_pt, assume_a="pos")
+
+        assert unify(pattern, out1) == {A: A_pt, b: b_pt}
+        assert unify(pattern, out2) is False
+
+        assume_a = var("assume_a")
+        pattern = etuple(
+            OpInstance(Blockwise, core_op=OpInstance(Solve, assume_a=assume_a)),
+            A,
+            b,
+        )
+        assert unify(pattern, out1) == {A: A_pt, b: b_pt, assume_a: "gen"}
+        assert unify(pattern, out2) == {A: A_pt, b: b_pt, assume_a: "pos"}
+
+
+    """
+
+    def __init__(
+        self,
+        op_type: type[Op],
+        parameters: dict[str, Any] | Sequence[tuple[str, Any]] | None = None,
+        **kwargs,
+    ):
+        if not (isinstance(op_type, type) and issubclass(op_type, Op)):
+            raise TypeError(f"Invalid op_type {op_type}. Expected type(Op)")
+
+        if kwargs:
+            if parameters is not None:
+                raise ValueError(
+                    "Cannot provide both parameters dict and keyword arguments"
+                )
+            parameters = kwargs
+        if isinstance(parameters, dict):
+            parameters = tuple(sorted(parameters.items()))
+        elif isinstance(parameters, list | tuple):
+            parameters = tuple(sorted(parameters))
+        elif parameters is None:
+            parameters = ()
+        self.op_type = op_type
+        self.parameters = parameters
+
+    def __str__(self):
+        return f"{self.op_type.__name__}({self.op_type}, {', '.join(f'{k}={v}' for k, v in self.parameters)})"
+
+
+def _unify_parametrized_op(v: Op, u: OpInstance, s: Mapping):
+    if not isinstance(v, u.op_type):
+        yield False
+        return
+    for parameter_key, parameter_pattern in u.parameters:
+        parameter_value = getattr(v, parameter_key)
+        s = yield _unify(parameter_value, parameter_pattern, s)
+        if s is False:
+            yield False
+            return
+    yield s
+
+
+_unify.add((Op, OpInstance, Mapping), _unify_parametrized_op)
+
+
 def convert_strs_to_vars(
     x: tuple | str | dict, var_map: dict[str, Var] | None = None
 ) -> ExpressionTuple | Var:
@@ -266,11 +365,13 @@ def convert_strs_to_vars(
     if var_map is None:
         var_map = {}
 
-    def _convert(y):
+    def _convert(y, op_prop=False):
         if isinstance(y, str):
             v = var_map.get(y, var(y))
             var_map[y] = v
             return v
+        if isinstance(y, LiteralString):
+            return y.value
         elif isinstance(y, dict):
             pattern = y["pattern"]
             if not isinstance(pattern, str):
@@ -282,8 +383,14 @@ def _convert(y):
             var_map[pattern] = v
             return v
         elif isinstance(y, tuple):
-            return etuple(*(_convert(e) for e in y))
-        elif isinstance(y, Number | np.ndarray):
+            return etuple(*(_convert(e, op_prop=op_prop) for e in y))
+        elif isinstance(y, OpInstance):
+            return OpInstance(
+                y.op_type,
+                {k: _convert(v, op_prop=True) for k, v in y.parameters},
+            )
+        elif (not op_prop) and isinstance(y, Number | np.ndarray):
+            # If we are converting an Op property, we don't want to convert numbers to PyTensor constants
             from pytensor.tensor import as_tensor_variable
 
             return as_tensor_variable(y)