Repository files navigation

Orbax provides common checkpointing and persistence utilities for JAX users.

Refer to our full documentationhere.

pip install orbax-checkpoint (latest PyPi release) OR

pip install 'git+https://github.com/google/orbax/#subdirectory=checkpoint'(from this repository, at HEAD)

import orbax.checkpoint

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

pip install orbax-export (latest PyPi release) OR

pip install 'git+https://github.com/google/orbax/#subdirectory=export'(from this repository, at HEAD)

import orbax.export

Orbax also includes a serialization library for JAX users, enabling theexporting of JAX models to the TensorFlow SavedModel format.

Note thatorbax-export requires TensorFlow, but does not include it by defaultto allow for flexibility in version choice. If you wish to install with standardTensorFlow, please usepip install orbax-export[all].

You can find the latest releases of Orbax Checkpointing atorbax-checkpoint and for Orbax Export atorbax-export.

Contactorbax-dev@google.com for help or with any questions about Orbax!

Orbax was initially published as a catch-all package itself. In order tominimize dependency bloat for users, we have frozen that package atorbax-0.1.6, and will continue to release future changes under thedomain-specific utilities detailed above (e.g.orbax-checkpoint).

As we have preserved the orbax namespace, existing import statements can remainunchanged (e.g.from orbax import checkpoint).