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:
- 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.
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.
2 maintainers
Metadata
2 maintainers
Release files for coordax 0.2.9
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| coordax-0.2.9.tar.gz | 69.0 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| coordax-0.2.9-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 145.0 kB
Release files / coordax-0.2.9.tar.gz
| Download URL | coordax-0.2.9.tar.gz |
|---|---|
| Size | 69.0 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
a2db8a9d17b95c239c47a34525d4fda912b20a0301f8b49c5fb07c6081c83896
|
|
BLAKE2b-256 checksum How to use checksums |
a99f65da256aaf3485fbd6bc58f297a3652ff28e26ff5302d235a8c411728d48
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.12.15
|
Release files / coordax-0.2.9-py3-none-any.whl
| Download URL | coordax-0.2.9-py3-none-any.whl |
|---|---|
| Size | 76.0 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
a62a5d3e190398e689a636dd3b3571fbf49fdf766be1d28d37f15f451a408284
|
|
BLAKE2b-256 checksum How to use checksums |
d3ab043ec5d899acd021a324bdc17c36149e8c576cf6b9433575fc06a4c96bab
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.12.15
|
