google/orbax

What is it?

Key features

• 01Save and restore JAX pytree state with a simple save/load API.
• 02Supports asynchronous checkpointing.
• 03Handles standard and custom types with flexible storage formats.
• 04Separate PyPI packages for checkpointing and exporting.

Use cases

• Saving model weights and optimizer state during JAX training runs.
• Restoring a training job from a previous checkpoint.
• Building a custom checkpointing workflow on top of JAX with composable components.

Quick start

1. 1

   Install the checkpointing package from PyPI

   pip install orbax-checkpoint
2. 2

   Save and restore a JAX pytree with the quickstart code from the README

Things to know

• !Requires JAX, so it is only useful if you already work with JAX models.

Topics

FAQ

What is orbax?

Checkpointing and persistence utilities for saving and restoring JAX model state, such as weights and optimizer state. Orbax is a Python library that provides common checkpointing and persistence utilities for JAX users. It handles saving and restoring pytree state such as model weights and optimizer state, so training code does not need to implement its own serialization logic.

Is orbax free to use? What license does it have?

orbax is open source under the Apache-2.0 license. Check the license terms before using it in commercial products.

Is orbax still maintained?

Actively maintained. The last commit to orbax was on Nov 7, 2026.

What can I use orbax for?

Saving model weights and optimizer state during JAX training runs. Restoring a training job from a previous checkpoint. Building a custom checkpointing workflow on top of JAX with composable components.

How do I install orbax?

Install the checkpointing package from PyPI: pip install orbax-checkpoint Save and restore a JAX pytree with the quickstart code from the README

Where can I get orbax?

The source code and installation instructions are on GitHub: https://github.com/google/orbax