orbax-checkpoint (pypi): Orbax Checkpoint
Keywords: JAX machine learning|checkpoint|training