Skip to content

Conversation

ricardoV94
Copy link
Member

@ricardoV94 ricardoV94 commented Aug 25, 2025

This allows matching multiple patterns to a single graph, avoiding repeated comparisons across similar patterns, which should scale nicely with the number of rewrites.

Also, unlike the current unification approach partial matching of Op properties can be done, which is nice as some Ops have many parameters that don't matter for rewrite purposes (such as inplace patterns). PatternSub couldn't handle this logic before.

Finally, this approach should be extendable to handle variadic/associative ops, similar to https://github.com/HPAC/matchpy


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

Copy link

codecov bot commented Aug 25, 2025

Codecov Report

❌ Patch coverage is 79.76190% with 34 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.53%. Comparing base (40ccab1) to head (dfef95b).

Files with missing lines Patch % Lines
pytensor/graph/rewriting/trie_unification.py 79.26% 21 Missing and 13 partials ⚠️

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

Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1592      +/-   ##
==========================================
- Coverage   81.54%   81.53%   -0.01%     
==========================================
  Files         230      231       +1     
  Lines       53136    53304     +168     
  Branches     9448     9496      +48     
==========================================
+ Hits        43329    43463     +134     
- Misses       7370     7391      +21     
- Partials     2437     2450      +13     
Files with missing lines Coverage Δ
pytensor/tensor/elemwise.py 90.07% <100.00%> (+0.05%) ⬆️
pytensor/graph/rewriting/trie_unification.py 79.26% <79.26%> (ø)
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@ricardoV94 ricardoV94 force-pushed the trie_unification branch 5 times, most recently from 229c14d to 0beabbb Compare August 26, 2025 14:01
@ricardoV94
Copy link
Member Author

On further investigation, this approach may not buy us much in PyTensor.

We got a nice idea spinned of to #1594, and if we think about it the tracks functionality works a bit like a first level matching of this trie functionality. Since our rewrites aren't very deep that's probably good enough, as we won't have many rewrites with long chains of similar checks (that are not already handled in the same dispatched function anyway).

This approach may still be useful if we move towards more symbolic outputs / integrations with egglog, but for now may be overkill.

It could be a chicken and egg situation, in that the reason we don't have many isolated long-chain rewrites (like similar variations of graphs, where you have an op before the other, or the other way around), is that we didn't have an efficient way to navigate through them like a trie-based approach would offer... So this PR can stay as a POC for future reference.

Also one may think about taking a step further and try to build an Aho-Corasick system to efficiently apply a Trie of rewrites over a graph without having to restart-from zero at every step. Again this is likely only useful if get a longer collection of deep rewrites.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant