pymc-devs / pymc-devs/pytensor
Add/Mul fusion does not reason about broadcasting
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
Extreme case, where we end up nearly duplicating runtime (because we add a + b in every iteration, instead of only once)
import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.compile.mode import Mode
a = pt.tensor("a", shape=(2,))
b = pt.tensor("b", shape=(2,))
c = pt.tensor("c", shape=(100_000, 2))
out = (a + b) + c # ideal associativity
fn1 = pytensor.function([a, b, c], out, mode="numba", trust_input=True)
with pytensor.config.change_flags(optimizer_verbose=True):
fn2 = pytensor.function([a, b, c], out, mode=Mode(linker="numba", optimizer=None), trust_input=True)
a_test = np.ones(a.type.shape)
b_test = np.ones(b.type.shape)
c_test = np.ones(c.type.shape)
np.testing.assert_allclose(fn1(a_test, b_test, c_test), fn2(a_test, b_test, c_test))
%timeit fn1(a_test, b_test, c_test) # 238 μs ± 14.8 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
%timeit fn2(a_test, b_test, c_test) # 154 μs ± 8.95 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
This happens in a couple places, AddCanonizer and flatten_nested_add_mul
We are careful not to do this in the regular Fusion. It may make sense to still canonicalize as variadic, but we may want to specialize into subsets that reduce the number of flops
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 AddCanonizer and flatten_nested_add_mul entry points mentioned in the issue, then run the provided broadcasting example in both optimized and unoptimized modes. Compare the generated work for (a + b) + c, including the regular Fusion behavior, and consider the canonicalization trade-offs described. Done means broadcasting cases avoid repeated additions or multiplications while preserving equivalent results.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- numpy, python
- Domain
- performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100