replaces numpy sqrt method with pytensor equivalent #6405
Merged
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
|
Would be good to check for other constants in other transforms/logp/logcdf methods |
Codecov Report
Additional details and impacted files@@ Coverage Diff @@
## main #6405 +/- ##
==========================================
+ Coverage 86.05% 94.74% +8.68%
==========================================
Files 148 148
Lines 27645 27645
==========================================
+ Hits 23791 26193 +2402
+ Misses 3854 1452 -2402
|
Do you need me to do this to merge @ricardoV94? Seems like it could take some time to go through all of the places where a numpy function is used.. |
|
No, it doesn't need to be done now, but it would be great 😃 |
|
We can certainly handle that in other PRs. |
|
Thanks for the contribution @morganstrom! |
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.
What is this PR about?
When using JAX on a Google hosted VM with TPU support, the float64 dtype is not supported. Unfortunately, numpy uses this as default and aggressively casts arrays to float64. This caused an issue where my model wouldn't compile at all. After debugging, I located the issue to the zerosum-normal transformation, where
np.sqrtis used. In this PR I'm replacing numpy with the pytensor equivalent, which fixes the issue.Checklist
Major / Breaking Changes
New features
Bugfixes
zerosumnormal's numpy operators with their pytensor counterparts to make them run on JAXDocumentation
Maintenance