Terrax is a Python package for building, training, and running AI-first Earth System Models (ESMs) and their components. It is built on jax, nnx, coordax, and xarray, enabling efficient hardware acceleration, powerful automatic differentiation, and seamless integration with the open-source geospatial ecosystem.
See CONTRIBUTING.md for details.
Apache 2.0; see LICENSE for details.
This project is not an official Google project. It is not supported by Google and Google specifically disclaims all warranties as to its quality, merchantability, or fitness for a particular purpose.
