Orbax: the checkpointing layer that JAX frameworks standardise on

Two packages, not one, and why that matters

Orbax provides common checkpointing and persistence utilities for JAX users, and the packaging is the first thing to understand.

The project is published on PyPI as separate domain-specific packages rather than one. The checkpointing package installs like this.

pip install orbax-checkpoint

And the exporting package is a separate install.

pip install orbax-export

Both can also be installed from GitHub at HEAD, using the subdirectory syntax that pip provides for monorepos.

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

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

The repository layout confirms the split. There are top-level checkpoint/, export/ and docs/ directories, alongside the Bazel build files WORKSPACE, BUILD and .bazelrc, a pyproject.toml, a readthedocs.yml and a .pylintrc.

That root pyproject.toml is worth a second look because it explains the monorepo convention. Its own comment says the distribution metadata, including project name, dependencies and build backend, lives in the per-package pyproject.toml files under checkpoint/, model/ and export/, and that only formatter and linter settings live at the root so they apply uniformly across the monorepo. The root file therefore contains only an isort configuration with the google profile and orbax declared as a known third party, and nothing else.

So a third package directory, model/, exists in the structure that the pyproject comment names, even though the readme's installation section only describes two packages. The readme does not document a model package, and that gap is worth noting rather than filling with guesses.

The primary language is Python, the licence is Apache-2.0, the default branch is main, and the last push was on 2026-10-20. The recent releases are v0.12.7 on 2026-10-06, v0.12.6 on 2026-09-26 and v0.12.5 on 2026-09-21, and the documentation is at orbax.readthedocs.io.

The quickstart, and what a checkpoint actually holds

The quickstart is short enough that the API shape is obvious from it.

import jax
from orbax.checkpoint import v1 as ocp

# Define your pytree state (e.g. weights, optimizer state)
state = {'a': jax.numpy.ones(2), 'b': 42}

# Save the state
ocp.save('/tmp/my_checkpoint', state)

# Restore the state
restored_state = ocp.load('/tmp/my_checkpoint')

Two things in that snippet carry the design.

First, the import is from orbax.checkpoint import v1 as ocp, not from orbax.checkpoint directly. The version namespace is explicit, which is the pattern a library uses when it is prepared to change its format and wants old checkpoints to keep working. The alias as ocp is then the ordinary way to keep the call sites short.

Second, the state is a pytree. The example is a dict with a jax.numpy array and a plain integer, and the comment says to define your pytree state, for example weights and optimizer state. That means the unit of checkpointing is a pytree, not an array, and the library works with the JAX data model rather than around it. A dict of arrays and a dict of a nested dict of arrays are the same kind of thing to Orbax.

That is the mechanism that makes it usable across frameworks. Flax, PaxML and T5X all have different parameter structures, and a checkpointing layer that speaks pytree works for all of them without knowing anything about their internal layouts.

The save and load calls are symmetric and take a directory path, so the checkpoint is a directory on disk rather than a single file. That is what allows the pluggable storage formats and the asynchronous behaviour: the format of what is inside the directory can vary, and it can be written by several processes at once.

What the project says it provides

The readme describes the scope in one sentence and then states the goal.

Orbax includes a checkpointing library oriented towards JAX users, supporting a variety of different features required by different frameworks, including asynchronous checkpointing, standard and custom types, and pluggable storage formats.

Asynchronous checkpointing is the feature that changes how a training loop is written. With a synchronous save, the training step blocks while the checkpoint is written, and for a large model that pause can be long enough to matter for throughput. With an asynchronous save, the training continues while the write proceeds, and the cost moves to the infrastructure, since the write needs somewhere to go and something to coordinate it. The readme names the capability but does not describe the API for it in the quickstart, so the documentation site is where to look.

Standard and custom types covers the other half of the problem. JAX pytrees can hold anything, and the built-in ones, arrays, numpy arrays, dictionaries and so on, are the standard types. Custom types are the ones a framework defines, and supporting them is what allows a framework's own parameter class to be checkpointed without being flattened into a dictionary first and reassembled afterwards by hand.

Pluggable storage formats is the third. The checkpoint is a directory, and what goes in it, and how it is laid out, is not fixed. That matters for two reasons: a format that suits a small array of parameters is not the same as one that suits a sharded multi-host training state, and a format you can change is a format you can migrate later.

The stated goal follows from those three: to provide a highly customizable and composable API that maximizes adaptability for diverse use cases.

Adaptable and customizable is a real strength when you are building a framework, and it is a real cost when you are running a job and just want a checkpoint. There is no single default described here that you cannot change, and no single answer to what the directory contains, which is why the readme keeps pointing at the full documentation rather than trying to summarise the API.

The adoption list, and what it tells you about the design

The readme has a section on existing users, and it is the most informative part of the document about why the API looks the way it does.

Orbax Checkpointing is used extensively across JAX machine learning frameworks and model implementations. Among the Google projects it lists Flax, Google's neural network library for JAX, Gemma, the open foundation models from Google DeepMind, Kauldron, a Google Research training and evaluation framework, PaxML, Google's high-performance framework for training large-scale JAX models, T5X, Google's JAX framework for high-performance sequence models, MaxText, a high-performance scalable JAX LLM implementation, MaxDiffusion, a stable diffusion JAX training library optimized for Cloud TPUs, Tunix, a JAX-native library for LLM post-training, and numerous Google-internal ML frameworks.

Outside Google it lists AXLearn, Apple's high-performance deep learning library built on top of JAX, and openpi, robotics foundation models by Physical Intelligence.

That list is a stress test for the design, and the diversity within it is the point. Flax is a neural network library with its own parameter abstractions. PaxML and T5X are training frameworks at different scales with different checkpoint formats of their own. MaxText is aimed at very large models on TPU pods. openpi is robotics, with entirely different data. A checkpointing layer that serves all of those cannot be opinionated about the shape of a parameter tree, which is why pytree is the unit rather than a framework-specific object.

The list also tells you where the pressure comes from. MaxText and MaxDiffusion are about scale and about multi-host training, and those are exactly the situations where synchronous saves hurt and where the storage format has to accommodate sharded state. The asynchronous checkpointing and pluggable storage features exist because of those users, not in spite of them.

The support section names both channels: the issue tracker for bugs and support requests, and the [email protected] address for direct help or questions.

There is also a paper for citing, Orbax: Distributed Checkpointing with JAX, available on arXiv, with a long author list spanning Google and Google DeepMind researchers.

How this compares with the checkpointing code it replaces

The alternatives are worth naming, because most JAX projects started with their own checkpoint code and moved to Orbax.

The first alternative is the framework's own mechanism. Flax has a serialization path, and several training frameworks have one. The problem with that is not that it is bad, it is that it is per-framework, so a project using Flax with PaxML-style training loops ends up converting between formats, and every format migration becomes your problem rather than the framework's.

The second is writing the serialization directly with jax.numpy.savez and friends. That works for a simple case, a flat dict of arrays, and stops working the moment you have sharded state, an optimizer state with a schedule, or a custom class. It also leaves you to design the on-disk format, and then you own its backwards compatibility forever.

The third is a general-purpose serialization library such as pickle-based approaches or a distributed framework's own checkpoint layer. The difference in approach is that those are not built around the JAX data model, so a pytree has to be flattened before it can be saved and reassembled after, and the mapping between the tree and the flat representation has to be maintained outside the library.

Orbax's bet is that the layer belongs to the ecosystem rather than to the framework, so that a checkpoint written by one framework can be read by another, and so that the storage format can evolve in one place.

The limitation is the same bet seen from the other side. A general layer has to be general, which means it cannot be exactly right for your case, and the pytree-plus-directory model means a checkpoint is a directory of files in a format Orbax defines rather than something you can read with a text editor. If you need a checkpoint you can inspect with standard tools, or you need to interoperate with a non-JAX system that expects a specific layout, that is a cost.

The 0.x version line is the other thing to weigh. The releases are all 0.12.x, and the presence of an explicit v1 namespace import in the quickstart tells you the API has moved before and the project expects to move again. That is normal for a Google-internal-adjacent library that is standardising across teams, and it is worth budgeting a little migration effort.

Editorial conclusion

Orbax fits a JAX team that wants one checkpointing implementation shared across its training and export stack rather than per-framework code, and it does not fit a project outside JAX, since it is built around pytrees. Before adopting it, decide which package you need, orbax-checkpoint or orbax-export, since the repository publishes them separately and the import paths differ, and note that the v1 API is imported as orbax.checkpoint.v1 rather than from the package root.

Frequently asked questions

Official sources

1. google/orbax on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases