Implement `OpPattern` for more flexible `tracks` and `PatternNodeRewriter` #1594

ricardoV94 · 2025-08-27T17:29:59Z

Spinoff from #1592

It introduces a new OpPattern that can be used for tracks and in PatternNodeRewriter

You can now do:

@node_rewriter(tracks=[OpPattern(Blockwise, core_op=OpPattern(Solve))])`

which wasn't possible before because there are many instances of Solve, unless you defined all ahead of time and enumerated them all in the tracks (including all variations of on_error and b_ndim...).

I added OpPattern tracks to almost all Blockwise (and some old Elemwise). Running the tests just in tests.tensor.rewriting.test_linalg, I saw a drop in the number of times the local node rewriters were called from 70k to 26k, with only 4k new calls to OpPattern parameter checks.

The bulk of the remaining calls from mostly from isolated in2out rewrites, where we don't iterate multiple times and therefore there's not much caching/sharing of tracks we can do. This will be further improved in #1607

The new OpPattern can also be used to beef up the PatternNodeRewriter

I further allowed one to use a callable for the output expression (or further reject the rewrite), which is much more flexible.

Here is an example that was not possible before (specifically matching any CAReduce with axis=None) and binding scalar_op to be used in the output Elemwise:

from pytensor.graph.rewriting.basic import PatternNodeRewriter
from pytensor.graph.rewriting.unify import OpPattern
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(
    (OpPattern(CAReduce, scalar_op="scalar_op", axis=None), 
        (OpPattern(Join), "join_axis", "a", "b")),
    output_fn,
)

📚 Documentation preview 📚: https://pytensor--1594.org.readthedocs.build/en/1594/

ricardoV94 · 2025-09-07T16:39:00Z

pytensor/tensor/rewriting/math.py

-            if inp.dtype != ottype:
-                inp = cast(inp, ottype)
-            return [inp]
+for pair in (


This avoids triggering the rewrite too many times, we spent an unreasonable amount of time in this rewrite the way it was structured, and rarely ended up using it

This info should be put into a comment, since it's unusual (in our code base anyway) that we're using codegen

Defining rewrites in a loop is not what I would call codegen :D

Added a comment, it does not have the trigger word codegen in it

review-notebook-app · 2025-09-07T20:26:21Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

codecov · 2025-09-07T21:08:14Z

Codecov Report

❌ Patch coverage is 84.00000% with 40 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.65%. Comparing base (d894350) to head (e12fcd4).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/graph/rewriting/basic.py	72.41%	13 Missing and 11 partials ⚠️
pytensor/graph/rewriting/unify.py	82.25%	7 Missing and 4 partials ⚠️
pytensor/tensor/rewriting/linalg.py	87.50%	3 Missing and 1 partial ⚠️
pytensor/scalar/basic.py	80.00%	0 Missing and 1 partial ⚠️

❌ Your patch check has failed because the patch coverage (84.00%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1594      +/-   ##
==========================================
- Coverage   81.70%   81.65%   -0.05%     
==========================================
  Files         230      230              
  Lines       52950    53002      +52     
  Branches     9404     9414      +10     
==========================================
+ Hits        43262    43280      +18     
- Misses       7256     7274      +18     
- Partials     2432     2448      +16

Files with missing lines	Coverage Δ
pytensor/graph/rewriting/kanren.py	`100.00% <100.00%> (ø)`
pytensor/tensor/_linalg/solve/rewriting.py	`95.34% <100.00%> (+0.03%)`	⬆️
pytensor/tensor/rewriting/basic.py	`95.70% <100.00%> (-0.01%)`	⬇️
pytensor/tensor/rewriting/blockwise.py	`96.29% <100.00%> (+0.05%)`	⬆️
pytensor/tensor/rewriting/elemwise.py	`92.92% <100.00%> (-0.45%)`	⬇️
pytensor/tensor/rewriting/math.py	`90.66% <100.00%> (+<0.01%)`	⬆️
pytensor/tensor/rewriting/subtensor.py	`91.07% <100.00%> (+0.07%)`	⬆️
pytensor/scalar/basic.py	`80.58% <80.00%> (+0.01%)`	⬆️
pytensor/tensor/rewriting/linalg.py	`91.09% <87.50%> (-0.98%)`	⬇️
pytensor/graph/rewriting/unify.py	`94.73% <82.25%> (-5.27%)`	⬇️
... and 1 more

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

ricardoV94 · 2025-09-07T21:11:03Z

Only mypy missing

jessegrabowski

It might be nice to be able to match on Op properties as well. Do you think that's in-scope for this first pass, or too complex?

Thinking specifically about dimshuffle, being able to do something like OpPattern(Dimshuffle, WithProperty('is_transpose', True))

pytensor/graph/rewriting/basic.py

jessegrabowski · 2025-09-08T04:05:12Z

pytensor/scalar/basic.py

@@ -2696,7 +2698,7 @@ class Sign(UnaryScalarOp):
    nfunc_spec = ("sign", 1, 1)

    @staticmethod
-    def output_types_preference(x):
+    def _output_types_preference(x):


Suggested change

def _output_types_preference(x):

def _get_output_types_preference(x):

Since it's a method in this case, not a member variable

This is what init of the base class will look for if I don't pass a function

I don't get it

Check for ScalarLoop.__init__, this plays a role there, and the name matters

pytensor/tensor/rewriting/basic.py

pytensor/tensor/rewriting/blockwise.py

pytensor/tensor/rewriting/linalg.py

jessegrabowski · 2025-09-08T04:11:27Z

pytensor/tensor/rewriting/linalg.py

+        # N.B. this can be further reduced to a yet-unwritten cho_solve Op
+        #     __if__ no other Op makes use of the L matrix during the
+        #     stabilization


I know you only changed the indentation level and thus don't care, but this Op exists now, so this note is out of date.

Should this rewrite go directly to cho_solve then?

ricardoV94 · 2025-09-08T05:41:33Z

Thinking specifically about dimshuffle, being able to do something like OpPattern(Dimshuffle, WithProperty('is_transpose', True)

You don't need a special object, it already works if Dimshuffle has that property/attribute. It uses getattr(op, param_name)

ricardoV94 · 2025-09-08T05:54:30Z

I can't reply to one of your comments so here it is:

It's still considered bad form to create a new member variable outside of __init__. Why isn't output_types_preference set to None as the default in the superclass constructor?

This API is very strange but I don't want to change it in this PR. They wanted this to be easy to change without creating new Op types (it's used for inplace for instance). Only change I did is the attribute always exists

Copilot

Pull Request Overview

This PR introduces a flexible OpPattern class for use in tracks and PatternNodeRewriter, enabling more powerful pattern matching in graph rewrites. Key improvements include replacing many hardcoded Op type checks with OpPattern-based tracks and adding support for parameterized Op matching.

Implements OpPattern class for flexible Op parameter matching in unification and tracks
Refactors existing node rewriters to use OpPattern tracks, reducing redundant code and improving performance
Enhances PatternNodeRewriter to support callable output patterns and OpPattern input patterns

Reviewed Changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
pytensor/graph/rewriting/unify.py	Implements OpPattern class and LiteralString for flexible Op matching
pytensor/graph/rewriting/basic.py	Enhances node rewriter infrastructure to support OpPattern tracks
pytensor/tensor/rewriting/basic.py	Adds elemwise_of helper and refactors some node rewriters to use OpPattern
pytensor/tensor/rewriting/blockwise.py	Adds blockwise_of helper for creating Blockwise OpPatterns
pytensor/tensor/rewriting/math.py	Refactors inverse function pairs using PatternNodeRewriter
pytensor/tensor/rewriting/linalg.py	Updates linear algebra rewrites to use blockwise_of OpPattern tracks
pytensor/tensor/rewriting/elemwise.py	Converts elemwise rewrites to use elemwise_of OpPattern tracks
pytensor/tensor/rewriting/subtensor.py	Updates subtensor rewrites to use blockwise_of patterns
pytensor/tensor/elemwise.py	Adds is_matrix_transpose property to DimShuffle
pytensor/scalar/basic.py	Fixes Sign output_types_preference implementation
tests/graph/rewriting/test_unify.py	Adds tests for OpPattern unification
tests/graph/rewriting/test_basic.py	Adds tests for OpPattern in PatternNodeRewriter

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

pytensor/tensor/rewriting/linalg.py

Also avoid repeated checks when an outer rewriter enforces tracks before calling individual node rewriters

ricardoV94 · 2025-09-11T10:26:00Z

I've told mypy to turn a blind eye enough that tests are passing :)

ricardoV94 added enhancement New feature or request graph rewriting labels Aug 27, 2025

ricardoV94 force-pushed the unify_op_instances branch from c103e75 to 9e89eca Compare August 28, 2025 17:46

ricardoV94 changed the title ~~Unify op instances~~ Implement OpPattern for more flexible tracks and PatternNodeRewriter Sep 6, 2025

ricardoV94 force-pushed the unify_op_instances branch from 9e89eca to ea37006 Compare September 7, 2025 16:28

ricardoV94 commented Sep 7, 2025

View reviewed changes

ricardoV94 force-pushed the unify_op_instances branch from ea37006 to b5fb239 Compare September 7, 2025 17:22

ricardoV94 mentioned this pull request Sep 7, 2025

Implement trie algorithm pattern matching #1592

Closed

ricardoV94 force-pushed the unify_op_instances branch from b5fb239 to 1d13782 Compare September 7, 2025 20:26

ricardoV94 force-pushed the unify_op_instances branch 2 times, most recently from f28fb72 to dd8cee9 Compare September 7, 2025 20:43

ricardoV94 marked this pull request as ready for review September 7, 2025 21:09

ricardoV94 requested review from Copilot and jessegrabowski September 7, 2025 21:10

This comment was marked as outdated.

Sign in to view

jessegrabowski requested changes Sep 8, 2025

View reviewed changes

Show errorcode in run_mypy

8d10764

ricardoV94 force-pushed the unify_op_instances branch from dd8cee9 to f76771a Compare September 10, 2025 10:47

ricardoV94 requested review from jessegrabowski and Copilot September 10, 2025 10:48

Copilot AI reviewed Sep 10, 2025

View reviewed changes

pytensor/tensor/rewriting/linalg.py Outdated Show resolved Hide resolved

ricardoV94 force-pushed the unify_op_instances branch 3 times, most recently from fe64030 to 021017c Compare September 11, 2025 10:08

Allow unifying with OpPattern

08cfb61

ricardoV94 added 2 commits September 11, 2025 12:18

Allow OpPattern in tracks

7e1cf63

Also avoid repeated checks when an outer rewriter enforces tracks before calling individual node rewriters

Use OpPattern in tracks

e12fcd4

ricardoV94 force-pushed the unify_op_instances branch from 021017c to e12fcd4 Compare September 11, 2025 10:18

	def _output_types_preference(x):
	def _get_output_types_preference(x):

Implement OpPattern for more flexible tracks and PatternNodeRewriter #1594

Are you sure you want to change the base?

Implement OpPattern for more flexible tracks and PatternNodeRewriter #1594

Conversation

ricardoV94 commented Aug 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 Sep 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

review-notebook-app bot commented Sep 7, 2025

Uh oh!

codecov bot commented Sep 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ricardoV94 commented Sep 7, 2025

Uh oh!

This comment was marked as outdated.

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Sep 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Sep 8, 2025

Uh oh!

ricardoV94 commented Sep 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

ricardoV94 commented Sep 11, 2025

Uh oh!

Uh oh!

Implement `OpPattern` for more flexible `tracks` and `PatternNodeRewriter` #1594

Implement `OpPattern` for more flexible `tracks` and `PatternNodeRewriter` #1594

ricardoV94 commented Aug 27, 2025 •

edited

Loading

ricardoV94 Sep 7, 2025 •

edited

Loading

codecov bot commented Sep 7, 2025 •

edited

Loading

ricardoV94 Sep 9, 2025 •

edited

Loading

ricardoV94 commented Sep 8, 2025 •

edited

Loading