pymc-devs / pymc-devs/pytensor

Should Alloc be pushed downstream of expand_dims

Open
#884 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

graph rewriting needs info performance
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.