Speedup graph traversal functions #1596
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
ricardoV94
force-pushed
the
improve_rewrites
branch
from
dbb4551 to
dcbcbd9
Compare
ricardoV94
changed the title
Codecov Report❌ Patch coverage is ❌ Your patch check has failed because the patch coverage (90.34%) 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 #1596 +/- ##
==========================================
- Coverage 81.71% 81.70% -0.02%
==========================================
Files 230 231 +1
Lines 52931 52922 -9
Branches 9403 9384 -19
==========================================
- Hits 43255 43238 -17
- Misses 7245 7252 +7
- Partials 2431 2432 +1
🚀 New features to boost your workflow:
|
ricardoV94
force-pushed
the
improve_rewrites
branch
2 times, most recently
from
20d88e5 to
c42ba61
Compare
ricardoV94
changed the title
ricardoV94
commented
Aug 29, 2025
ricardoV94
force-pushed
the
improve_rewrites
branch
from
c42ba61 to
27ede5a
Compare
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
ricardoV94
force-pushed
the
improve_rewrites
branch
from
27ede5a to
262bd26
Compare
ricardoV94
force-pushed
the
improve_rewrites
branch
8 times, most recently
from
3666a4b to
afdfd7b
Compare
ricardoV94
changed the title
ricardoV94
force-pushed
the
improve_rewrites
branch
7 times, most recently
from
deeba35 to
a2d80ed
Compare
ricardoV94
force-pushed
the
improve_rewrites
branch
from
a2d80ed to
5ec8e59
Compare
ricardoV94
requested review from
Copilot and
jessegrabowski
and removed request for
Copilot
There was a problem hiding this comment.
Pull Request Overview
This PR optimizes graph traversal and topological sorting functions across PyTensor to improve performance, especially for rewrite operations. The main changes involve moving graph traversal functionality to a dedicated module and implementing more efficient algorithms.
Key changes:
- Rewrote traversal functions for 3x faster ancestors iteration and 2x faster topological sorting
- Converted methods to generators to reduce memory allocation
- Added a new "dfs" order to WalkingGraphRewriter for 5-6x iteration speedup
- Moved traversal functionality from
pytensor.graph.basictopytensor.graph.traversal
Reviewed Changes
Copilot reviewed 65 out of 65 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
pytensor/graph/traversal.py |
New module containing optimized graph traversal algorithms and functions |
pytensor/graph/basic.py |
Removed traversal functions, retained core graph structures |
pytensor/graph/rewriting/basic.py |
Added "dfs" order option and updated imports for traversal functions |
| Multiple test files | Updated imports to use traversal functions from new module |
| Multiple rewriting modules | Changed from in2out/out2in to dfs_rewriter for better performance |
Comments suppressed due to low confidence (7)
pytensor/graph/fg.py:1
- The error message pattern has changed from 'not reversible' to 'not iterable', but this might indicate a behavioral change in how the validation works. Verify that this change in error message is intentional and reflects the actual validation logic.
"""A container for specifying and manipulating a graph with distinct inputs and outputs."""
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
ricardoV94
force-pushed
the
improve_rewrites
branch
from
5ec8e59 to
276e5a9
Compare
* Avoid reversing inputs as we traverse graph * Simplify io_toposort without ordering (and refactor into its own function) * Removes client side-effect on previous toposort functions * Remove duplicated logic across methods
ricardoV94
force-pushed
the
improve_rewrites
branch
from
276e5a9 to
7353874
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Graph traversal / toposort is used widely by our rewrites.
This PR rewrites most methods for performance, by ignoring the DRY principle and hoisting checks out of hot loops.
It further converts existing methods to generators, to avoid alocating large intermediate memory, when iterating one item at a time suffices, or when we end up copying immediately to another container (like a deque)
MAJOR CHANGE: We avoid reversing the inputs as we iterate, which gives a major change in the order of the returned variables. Instead of doing a left-recursive depth-first search we end up doing a right-recursive depth-first search. This may be a bit confusing, in that stuff like
ancestors(pt.exp(x, y))will yieldy, xin that order.Hopefully people aren't relying on the order of the inputs when they retrieve them with these helpers.
But if there is a need for it, we can easily add a
left_recursiveboolean flag toancestorsandtoposort(and every function that uses those), to retrieve the old behavior. Note thatwalkandwalk_toposortare completely flexible in the order of iteration since the expansion comes from a user provided function (when we are the ones defining the functions, we could also add the option to reverse during expansion).Iterating over the ancestors is now ~3x faster and toposort ~2x faster than before. Furthermore, several rewrites don't really need to be applied in topological order, so I added a new
"dfs"order to theWalkingGraphRewriterand used that instead. This should speedup ~5-6x iteration speed over the nodes (besides what's saved by avoiding the full list allocation / copy).Some benchmarks
This PR decides to move the traversal functionality to its own file, which I think is more manageable. There are some backward compatible imports with FutureWarnings.
Another breaking change is that the
io_toposortandgeneral_toposortno longer accept the special keyword arguments.clientsinio_toposortandgeneral_toposort, which had a single internal use in one scan rewrite that can be avoided by just requesting them from the innerFunctionGraph(this didn't use to exist back in the day).compute_deps_cacheandcompute_depsingeneral_toposort, whichio_toposortmade use of to reduce one level of callable nesting. This is really minor now that we cleaned up other things and it was likely also only an internal usage.Also remove the
.owneraccess indirection (behind a property), which adds extra cost for a very very common accessed attribute. You can see in the benchmarks above that the effect is visible there, but it should go beyond traversal functions.Also cleanup FunctionGraph methods to reduce attribute access, single line function or expensive checks.