GitHub repository: tsukuba-hpcs/orbax
Description: Orbax provides common checkpointing and persistence utilities for JAX users