Orbax is a library providing common utilities for JAX users.

Orbax includes a checkpointing library oriented towards JAX users, supporting a variety of different features required by different frameworks, including asynchronous checkpointing, various types, and various storage formats. We aim to provide a highly customizable and composable API which maximizes flexibility for diverse use cases.

To get started, check out our documentation.

Check out our colab for a hands-on introduction.