pymc-devs / pymc-devs/pytensor
Missing some simple matrix algebraic simplifications
Open
Nobody has claimed this yet.
graph rewriting
linalg
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
Missing:
A@B + A@C = A@(B+C)(one less matmul)s*A @ B = s*(A@B)(which can be done by a single gemm routine)
from pytensor.graph import rewrite_graph
import pytensor.tensor as pt
A,B,C = pt.matrices("ABC")
s = pt.scalar("s")
o1 = A@B + A@C
rewrite_graph(o1, include=("fast_run",), exclude=("inplace",)).dprint()
print()
o2 = (s*A) @ B
rewrite_graph(o2, include=("fast_run",), exclude=("inplace",)).dprint()
Gemm{no_inplace} [id A]
├─ Dot22 [id B]
│ ├─ A [id C]
│ └─ B [id D]
├─ 1.0 [id E]
├─ A [id C]
├─ C [id F]
└─ 1.0 [id E]
Dot22 [id A]
├─ Mul [id B]
│ ├─ ExpandDims{axes=[0, 1]} [id C]
│ │ └─ s [id D]
│ └─ A [id E]
└─ B [id F]
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
Reproduce both cases with the supplied rewrite_graph examples and inspect the fast_run rewrites for matrix operations. Done means the first expression is represented as one shared-left matrix multiplication with an addition, and the scaled multiplication uses a single GEMM-compatible operation rather than multiplying the input first.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, performance
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100