coordax 0.2.1

Coordinate axes for scientific computing in JAX

Coordax: Coordinate Axes for JAX

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:

1. First class integration with JAX, including support for arbitrary JAX transformations
2. Easy wrapping of code not written for labeled arrays with cmap, inspired by Penzai
3. Optional Coordinate objects, for advanced use-cases
4. 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.

Disclaimer

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.