Introducing Orbax, the new checkpointing library of Flax

[IvyZX]

Hi folks, we are happy to announce that Flax is starting the first step to migrate the checkpointing features to Orbax (go/orbax), starting from 0.6.3.

Orbax is a highly customizable and flexible library for saving various array types and storage formats, and is already adopted by a handful of Google JAX models like T5X.

You are encouraged to try out Orbax starting from today! Just install from head (or along with pip install flax) and add an import line and an extra arg to your typical save_checkpoint() or save_checkpoint_multiprocess() calls:

pip3 install -U "git+https://github.com/google/orbax.git@main#egg=orbax"

import orbax.checkpoint as orbax
save_checkpoint(..., orbax_checkpointer=orbax.Checkpointer(orbax.PyTreeCheckpointHandler()))

For more use cases (e.g. async saving) and code examples, please check out our official checkpointing guide. Or reach out to us if you have any questions. 😀 Happy checkpointing!