-
Notifications
You must be signed in to change notification settings - Fork 6
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
Use jaxtyping to enrich type annotations #5
Comments
I've also been following Observations on my end:
Let me know if you more thoughts on this, or bandwidth for contributing. 🙂 |
Surely the first point should be verified. |
Looking forward to it, thanks @lucagrementieri! |
From my experiments I think it's possible to support all the features using I think my pull request should be ready for next week. |
I finally had time to push my PR #6 ! Sorry for the delay! |
I just discovered the jaxtyping library and I think it could be an interesting alternative to the current typing system proposed by jax_dataclasses.
jaxtyping supports variable-size axes and symbolic expressions in terms of other variable-size axes, see https://github.com/google/jaxtyping/blob/main/API.md and it has very few requirements.
Do you think that it could be added to jax_dataclasses?
The text was updated successfully, but these errors were encountered: