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
-
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