pymc-devs / pymc-devs/pytensor
Don't let tensor creation Ops constant fold too soon
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
It's useful to know when arange or eye is in a graph, but if these have constant inputs they will be constant folded, and then we would need to analyze all constants in a graph to find whether they correspond to arange/eye, which can be rather costly to do.
We should prevent these from being constant folded until a later phase? Also we may want to do a single pass at the beginning that converts constants to equivalent arange, eye, alloc. As I mentioned this can be costly, but at least this way it's "cached" because it becomes part of the graph. Whereas a rewrite looking for aranges would otherwise have to do it for every constant every time it is called (or equally unsatisfactory, ignore constants).
This is actually one of the hardest aspects of the kind of destructive/eager optimization we do, as constant_fold destroys information but can also simplify graphs very quickly.
For reference, note that there already exists special logic is alloc to try to not constant_fold it in a few cases where it is worse not to. This is just a bit myopic and won't scale for the kind of rewrites we are interested in #573
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the alloc special logic in pytensor/tensor/basic.py at lines 1675-1713, then trace the constant-folding phase and handling of arange and eye. Determine how these tensor-creation operations can remain identifiable until a later phase or be converted in one initial pass; done means the graph preserves or caches this information without repeatedly scanning all constants.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100