Authors: Dmitrii Kochkov and Stephan Hoyer
Coordax is a Python library for labeled axes with JAX. Our approach is reminiscent of Xarray, but tailored to meet the needs of modern physics- and AI-based simulation codes written in JAX, such as NeuralGCM.
Compared to other libraries for labeled arrays, Coordax provides a handful of key features:
- First class integration with JAX, including support for arbitrary JAX transformations
- Easy wrapping of code not written for labeled arrays with
cmap, inspired by Penzai - Optional
Coordinateobjects, for advanced use-cases - Lossless conversion to and from Xarray, for serialization and data analysis
Coordax was developed to meet the needs of NeuralGCM, but we hope it will be useful more broadly!
For more details, read the documentation.
Coordax is an experiment that we are sharing with the outside world in the hope that it will be useful. It is not a supported Google product. We welcome feedback, bug reports and code contributions, but cannot guarantee they will be addressed.
