pymc-devs / pymc-devs/pytensor
An Elemwise feeding `dest[idx].set(...)` / `.inc(...)` isn't fused into the write (extra temp + copy)
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
When an Elemwise result is written into a subtensor via .set(...)/.inc(...), the Elemwise output is materialized in its own buffer and then copied into the destination, even when the write itself is in-place. The Elemwise could instead write directly into the destination region.
Repro (Numba backend; destination intermediate so the write is in-place):
import pytensor, pytensor.tensor as pt
x = pt.vector("x")
o0 = pt.vector("o0")
o = o0 + 1.0
res = o[1:].set(pt.exp(x))
fn = pytensor.function([x, o0], res, mode="NUMBA")
pytensor.dprint(fn, print_memory_map=True)
SetSubtensor{start:} d={0: [0]} # in-place on o
├─ Add (o0 + 1)
├─ Exp # own buffer (inplace_pattern={}, destroy_map={})
│ └─ x
└─ 1
So per call we allocate exp(x) in its own buffer and then copy it into o[1:]. The in-place SetSubtensor only avoids copying the whole o, not the Elemwise temp.
Elemwise inplace can only destroy one of its own inputs; the write destination o[1:] is not an input to Exp, so there's no path for the Elemwise to write its result directly into the destination.
The advanced-indexing fusion (IndexedElemwise) already avoids this (after #2015) — its inner graph is AdvancedIncSubtensor1{inplace,set}(buffer, Exp(...), idx), so the Elemwise is fused into the scatter and the result lands in the destination with no temp. The basic Subtensor/IncSubtensor path has no equivalent rewrite.
We should generalize the IndexedElemwise fuse-elemwise-into-write machinery to also absorb basic Subtensor/IncSubtensor, unifying the slice and advanced paths.
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 provided Numba reproduction and memory map to confirm the extra Elemwise buffer and copy. Then trace the IndexedElemwise fusion machinery described as following #2015 and compare it with the basic Subtensor/IncSubtensor path. Done means the basic slice write absorbs the Elemwise result so it lands directly in the destination without the temporary buffer or copy.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100