This library supports using xarray datatypes together with the JAX library.
This project has been spun off from the xarray_jax utility in the graphcast project.
We plan to make some improvements to it before migrating existing uses (including graphcast) to it. These improvements are not yet complete and so we don't advise migrating to this code yet.
We will aim to make improvements largely backwards-compatible with
graphcast.xarray_jax, but they may not be entirely so.