Fix shape issues in jax tridiagonal solve #1414
Open
+63
−1
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
jessegrabowski
changed the title
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #1414 +/- ##
==========================================
+ Coverage 82.12% 82.14% +0.01%
==========================================
Files 211 211
Lines 49687 49695 +8
Branches 8813 8815 +2
==========================================
+ Hits 40807 40821 +14
+ Misses 6702 6697 -5
+ Partials 2178 2177 -1
🚀 New features to boost your workflow:
|
jessegrabowski
requested changes
May 23, 2025
| dl = jax.numpy.pad(dl, (1, 0)) | ||
| du = jax.numpy.pad(du, (0, 1)) | ||
| # if b is a vector, broadcast it to be a matrix | ||
| b_is_vec = len(b.shape) == 1 |
There was a problem hiding this comment.
You don't need to check this at runtime. The Solve Op has a property b_ndim, so you can do:
b_is_vec = op.b_ndim
if assume_a == 'tridiagonal':
... # carry onThere was a problem hiding this comment.
The check as written will also fail in the batched case (that's why we have it at the Op level)
There was a problem hiding this comment.
you'd have to do b_is_vec = op.b_ndim == 1 though, no? because bool(op.b_ndim) -> True for b_ndim > 0
There was a problem hiding this comment.
yes exactly, my code has an error
Co-authored-by: Jesse Grabowski <48652735+jessegrabowski@users.noreply.github.com>
benmaier
commented
May 23, 2025
| ), | ||
| (4, 5, 5), | ||
| 2, | ||
| ), |
There was a problem hiding this comment.
This is what ruff came up with..
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
In the tridiagonal solve dispatch for jax, dl and du should have same shape as d, and b should have rank2. This PR:
dlandduso they have the same shape asdbis a vector -- if so, add a dimension that makes the vector a column vector, compute the result and cast back to vectorRelated Issue
Checklist
Type of change
📚 Documentation preview 📚: https://pytensor--1414.org.readthedocs.build/en/1414/