pymc-devs / pymc-devs/pytensor
Implement `linalg.BandedDot`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
Benchmarking in #1323 showed that the banded (tridiagonal) case can get huge speedups by using a specialized dot product. This issue asks for a BandedDot Op that uses xgbmv to realize these speedups.
In the future it would be nice to be able to rewrite into this Op in cases were we see that we can, but I don't think it's necessary on first pass. Just having the functionality laying around will be nice.
Note that JAX doesn't use xgbmv to do this in the tridiagonal case. They have _tridiagonal_product that just directly does it using jax primitive Ops. This might be preferable, because it would require no extra dispatch work, but it would not let us handle the general banded case -- only the tridiagonal case. Maybe we want both?
At minimum, we should benchmark xgbmv vs direct method in the tridiagonal case.
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 locating the linalg.BandedDot entry point and existing linear-algebra Ops. Benchmark xgbmv against the direct method for the tridiagonal case, then implement the BandedDot Op using xgbmv for the general banded case. Done means the Op is available and the requested performance comparison is recorded.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 32/100