pymc-devs / pymc-devs/pytensor
Optimize `Sum`s of `MakeVector`s and `Join`s
Open
Nobody has claimed this yet.
beginner friendly
graph rewriting
performance
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Please describe the purpose of filing this issue
Not a drastic improvement by any means, but something we can keep in mind:
reduce(at.concatenate(*tensors)) -> reduce(reduce(tensor) for tensor in tensors)
Ignoring any axis complexities
import pytensor
import pytensor.tensor as pt
import numpy as np
x = pt.vector("x")
y = pt.vector("y")
f1 = pytensor.function([x, y], pt.sum(pt.concatenate((x, y))))
f2 = pytensor.function([x, y], pt.sum((pt.sum(x), pt.sum(y))))
f3 = pytensor.function([x, y], pt.add(pt.sum(x), pt.sum(y)))
pytensor.dprint(f1)
print()
pytensor.dprint(f2)
print()
pytensor.dprint(f3)
x_val = np.random.rand(100_000)
y_val = np.random.rand(200_000)
%timeit f1(x_val, y_val)
%timeit f2(x_val, y_val)
%timeit f3(x_val, y_val)
Sum{acc_dtype=float64} [id A] '' 1
|Join [id B] '' 0
|TensorConstant{0} [id C]
|x [id D]
|y [id E]
Sum{acc_dtype=float64} [id A] '' 3
|MakeVector{dtype='float64'} [id B] '' 2
|Sum{acc_dtype=float64} [id C] '' 1
| |x [id D]
|Sum{acc_dtype=float64} [id E] '' 0
|y [id F]
Elemwise{Add}[(0, 0)] [id A] '' 2
|Sum{acc_dtype=float64} [id B] '' 1
| |x [id C]
|Sum{acc_dtype=float64} [id D] '' 0
|y [id E]
544 µs ± 27.5 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
270 µs ± 5.11 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
270 µs ± 8.86 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
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 reproducing the provided PyTensor script and inspecting the printed graphs for the three sum formulations. Investigate the existing graph-optimization entry points that handle Sum, MakeVector, and Join; done means the equivalent concatenation form is rewritten safely and the benchmark shows the intended improvement without changing results.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- numpy, python
- Domain
- backend, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100