pymc-devs / pymc-devs/pytensor
Should Alloc be pushed downstream of expand_dims
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
We have some other rewrites that will push Alloc below Elemwise, so that we don't compute on repeated inputs, but this won't happen if there's an expand_dims in the way. As of now the following graph does not get lifted
import pytensor
import pytensor.tensor as pt
x = pt.vector("x", shape=(3,))
y = pt.alloc(x, 1000, 3)[None]
out = pt.exp(y)
pytensor.function([x], out).dprint(print_type=True)
# Exp [id A] <Tensor3(float64, shape=(1, 1000, 3))> 1
# └─ Alloc [id B] <Tensor3(float64, shape=(1, 1000, 3))> 0
# ├─ x [id C] <Vector(float64, shape=(3,))>
# ├─ 1 [id D] <Scalar(int8, shape=())>
# ├─ 1000 [id E] <Scalar(int16, shape=())>
# └─ 3 [id F] <Scalar(int8, shape=())>
There is actually an "uncanonicalize" rewrite that allows "lifting" expand_dims above some Alloc, which would have helped here.
https://github.com/pymc-devs/pytensor/blob/e6e6d69f6d878786270f1751098b0682e2d8f607/pytensor/tensor/rewriting/uncanonicalize.py#L125-L150
However, this is at odds with the opposite canonical local_alloc_sink_dimshuffle:
https://github.com/pymc-devs/pytensor/blob/d62f4b19d412d91994dd12362f7976d690911084/pytensor/tensor/rewriting/basic.py#L462-L467
It's not obvious to me why the latter should be given preference. In general it seems like we can always lift expand_dims towards the inputs of the function (as it does not affect number of operations), and sink alloc towards the outputs. But here we are not allowing the "swap" when an expand_dims meets an alloc
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
Run the provided PyTensor example and inspect the rewrite interactions in pytensor/tensor/rewriting/uncanonicalize.py and pytensor/tensor/rewriting/basic.py, especially the linked rewrite sections. Determine the intended ordering for expand_dims and Alloc, then add a regression test showing the chosen canonical graph and confirming that repeated computation is avoided.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100