Orbax#

Orbax is an umbrella namespace providing common training utilities for JAXusers. It includes multiple distinct but interrelated libraries.

Checkpointing

A flexible and customizable API for managing checkpoints consistingof various user-defined objects in multi-host, multi-device settings.

Exporting

A library for exporting JAX models to Tensorflow SavedModel format.

Installation#

There is no singleorbax package, but rather a separate package for eachfunctionality provided by the Orbax namespace.

The latest release oforbax-checkpoint can be installed fromPyPI using

pipinstallorbax-checkpoint

You may also install directly from GitHub, using the following command. Thiscan be used to obtain the most recent version of Optax.

pipinstall'git+https://github.com/google/orbax/#subdirectory=checkpoint'

NOTE: Certain edge cases oforbax-checkpoint may not work on Windows.

Also, supporting them is not planned yet.

Similarly,orbax-export can be installed fromPyPI using

pipinstallorbax-export

Install from GitHub using the following.

pipinstall'git+https://github.com/google/orbax/#subdirectory=export'

Checkpointing#

Exporting#

Support#

Please report any issues or request support using ourissue tracker.