Xarray <-> Jax #124
allen-adastra
started this conversation in
Ideas
Xarray <-> Jax
#124
Replies: 1 comment
-
|
Update: new experimental repo that gets at least part of the way there: https://github.com/allen-adastra/xarray_jax |
Beta Was this translation helpful? Give feedback.
0 replies
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Uh oh!
There was an error while loading. Please reload this page.
-
It would be very helpful (for POPSIM as well) to have proper integration between Xarray + Jax. That way, we could have metadata + coordinates travelling around with our spatial variables.
There has been some discussion on this matter: jax-ml/jax#17107 and google-deepmind/graphcast#9
It seems the right time to do it is when both sides adopt support for the Array API. Relevant issues for Jax jax-ml/jax#18353 (which seems to be getting close) and Xarray pydata/xarray#7848
Not sure if there is something we can do here to help move things along :)
Beta Was this translation helpful? Give feedback.
All reactions