pymc-devs / pymc-devs/pytensor
numba max reduction is 10x slower than sum on the same array
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
pt.max reduces at about a tenth the speed of pt.sum through the same CAReduce codegen on the same contiguous array, and unlike sum it loses to numpy rather than beating it. Every axis configuration I tried, full and partial.
import time
import numpy as np
import pytensor
import pytensor.tensor as pt
x_np = np.random.default_rng(0).normal(size=(200, 200, 200))
x = pt.tensor("x", shape=x_np.shape)
def bench(fn, *args):
fn(*args)
start = time.perf_counter()
for _ in range(10):
fn(*args)
return (time.perf_counter() - start) / 10 * 1000
for name, op in (("max", pt.max), ("sum", pt.sum)):
print(name, round(bench(pytensor.function([x], op(x), trust_input=True), x_np), 2), "ms")
print("numpy max", round(bench(x_np.max), 2), "ms")
print("numpy sum", round(bench(x_np.sum), 2), "ms")
# max 7.31 ms sum 0.72 ms numpy max 0.72 ms numpy sum 0.96 ms
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 by running the Python benchmark in the issue, comparing pt.max and pt.sum through the CAReduce codegen on full and partial axes. Trace the reduction entry points and generated execution path to identify why max is slower; done means max approaches NumPy performance and no longer has the reported gap with sum across the tested axis configurations.
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
- 50/100