-
Notifications
You must be signed in to change notification settings - Fork 42
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.
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
initial working jaxtyping serializing/deserializing #576
Conversation
The code changes in `serde/numpy.py` modify the `is_numpy_jaxtyping` function to handle a `TypeError` exception. This change ensures that if a `TypeError` occurs when trying to determine if a type is a numpy jaxtyping, the function will return `False` instead of raising an error.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #576 +/- ##
==========================================
- Coverage 91.04% 90.86% -0.18%
==========================================
Files 13 13
Lines 1942 1971 +29
Branches 437 440 +3
==========================================
+ Hits 1768 1791 +23
- Misses 117 122 +5
- Partials 57 58 +1 ☔ View full report in Codecov by Sentry. |
Hi @yukinarit, I've implemented initial support for jaxtyping serialization/deserialization and all tests are passing. Before I proceed with adding specific tests for jaxtyping, could you please review the current implementation and let me know if you're happy with the direction? Any feedback would be greatly appreciated! |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The PR basically looks good to me.
But I am not a numpy/jaxtyping user. @kmsquire It would be great if you can take a look and give us feedback 🙏
Made jaxtyping an optional dep and addressed your typ comment! |
Added tests in! Let me know if there's anything else you'd like me to add/fix @yukinarit @kmsquire, otherwise I also think this is probably good to merge |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nice work! LGTM!
@pablovela5620 Is there anything you want to add before publishing v0.20.0? |
Nope, this looks good! Thanks |
Jaxtyping serialization/deserializing working along with type checking with beartype #572 . This allows one to use jaxtyping when serializing/deserializing and check both the dtype and the shape. Let me know if this direction makes sense!