Repository navigation
Conversation
|
This pull request has been imported. If you are a Meta employee, you can view this in D126197644. (Because this pull request was imported automatically, there will not be any future comments.) |
stroxler
left a comment
stroxler
left a comment
There was a problem hiding this comment.
This looks pretty good to me, although I think we need to work out why tests are failing and fix that.
cc @jakevdp - could you double-check this, in particular the basearray.pyi module. I want to make sure we don't introduce any masking bugs, where some names become unavailable. I have pretty limited availability this week due to Pytorch Conference.
|
Thanks – I have a similar fix underway, happy to use yours if you'd prefer. Note however that the CI testing is still being done with jax v0.11.2: that's specified here: pyrefly/tensor-shapes/test-requirements.txt Lines 22 to 25 in b02aca9 Edit: there are actually a bunch more typing changes in JAX v0.12 that are not reflected in this PR – I think I'll send mine anyway. |
|
One comment – I think it would be better to move the changes not associated with jax v0.12 to a separate PR |
|
My JAX v0.12 updates are in #5301. |
ae0a234 to
d949e3b
Compare
bdelwood
changed the title
Summary: Follows #5291. The remaining gaps I encountered were in numpy: scalar arithmetic and conversions, array arithmetic with complex and numpy scalars, and missing arguments to arange, round, min, and max. This PR adds support for those operations without attempting to fully model numpy’s dtype promotion rules. - [x] I am a human, and I fully understand the code I am submitting. Pull Request resolved: #5292 Test Plan: Added a few tests. Reviewed By: stroxler Differential Revision: D126199088 fbshipit-source-id: d7219192e7336e23e4991cbef5f79ef501ad8cca
stroxler
left a comment
stroxler
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
Summary: Builds on #5301. Split from #5291. This adds the lax.polynomial export introduced in jax 0.12 and a test covering its use. - [x] I am a human, and I fully understand the code I am submitting. Pull Request resolved: #5302 Reviewed By: stroxler Differential Revision: D126549442 fbshipit-source-id: 730d45131ef9836a77b1ae66ec4a92b65f978935
|
This pull request has been merged in c6e50c3. |
4 participants
Summary
Next after #5290
I found that these stubs did not always agree with jax’s own annotations, causing valid calls to be rejected. This PR makes the ones I ran into issues with consistent. It also adds support for numpy boolean scalars in arithmetic.
Test Plan
Added a couple of tests.