pymc-devs / pymc-devs/pytensor
Recognize `dot` from naive sum of broadcasted muls
Open
Nobody has claimed this yet.
graph rewriting
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
Brought up in https://github.com/pymc-devs/pytensor/pull/858
import pytensor
import pytensor.tensor as pt
a = pt.matrix("a", shape=(200, 300))
b = pt.matrix("b", shape=(300, 400))
dot = (a[:, :, None] * b).sum(1)
fn = pytensor.function([a, b], dot)
pytensor.dprint(fn, print_type=True)
# Sum{axis=1} [id A] <Matrix(float64, shape=(200, 400))> 3
# └─ Mul [id B] <Tensor3(float64, shape=(200, 300, 400))> 2
# ├─ ExpandDims{axis=2} [id C] <Tensor3(float64, shape=(200, 300, 1))> 1
# │ └─ a [id D] <Matrix(float64, shape=(200, 300))>
# └─ ExpandDims{axis=0} [id E] <Tensor3(float64, shape=(1, 300, 400))> 0
# └─ b [id F] <Matrix(float64, shape=(300, 400))>
fn_dot = pytensor.function([a, b], a @ b)
print(); pytensor.dprint(fn_dot, print_type=True)
# Dot22 [id A] <Matrix(float64, shape=(200, 400))> 0
# ├─ a [id B] <Matrix(float64, shape=(200, 300))>
# └─ b [id C] <Matrix(float64, shape=(300, 400))>
a_test = np.random.normal(size=a.type.shape)
b_test = np.random.normal(size=b.type.shape)
np.testing.assert_allclose(fn(a_test, b_test), fn_dot(a_test, b_test))
%timeit fn(a_test, b_test) # 70.9 ms ± 1.29 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit fn_dot(a_test, b_test) # 861 µs ± 148 µ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 running the issue's naive broadcasted multiplication and summation example through pytensor.function and pytensor.dprint, then compare it with a @ b. Trace the graph optimization entry points that handle this expression; done means the naive form is recognized as an equivalent dot operation and its outputs match the direct dot form.
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
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100