pymc-devs / pymc-devs/pytensor

Add/Mul fusion does not reason about broadcasting

Open
#1,966 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.