Skip to content

Implement linalg.BandedDot #1415

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
jessegrabowski opened this issue May 23, 2025 · 2 comments · May be fixed by #1416
Open

Implement linalg.BandedDot #1415

jessegrabowski opened this issue May 23, 2025 · 2 comments · May be fixed by #1416
Labels
feature request help wanted Extra attention is needed linalg Linear algebra Op implementation

Comments

@jessegrabowski
Copy link
Member

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.
@jessegrabowski
Copy link
Member Author

@benmaier tagging you because this is related to the PDE stuff

@jessegrabowski
Copy link
Member Author

I updated my benchmark gist to include the "direct tridiagonal" case that jax uses. It is much worse than xgbmv, but maybe that's not surprising? It's faster than A @ b on the very largest cases, but not otherwise.

@jessegrabowski jessegrabowski linked a pull request May 23, 2025 that will close this issue
Open
11 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
feature request help wanted Extra attention is needed linalg Linear algebra Op implementation
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

Implement BandedDot Op jessegrabowski/pytensor
1 participant
@jessegrabowski