Checkpointing and Exporting JAX Models: An End-to-End Guide with Orbax#
This guide demonstrates a complete, end-to-end workflow for managing JAX models using the Orbax library, from robust training-time checkpointing to final model export. We will simulate a Flax/Optax setup to show how the Checkpointer API enables policy-based management and restoration of training states. Following that, we use the standalone save function to save the final parameters for inference. At the end, we export these parameters into a TensorFlow SavedModel with orbax-export.
1. Setup#
First, we set up the necessary environment by installing the required packages and importing the modules we’ll use throughout this guide.
Note: The following cells install the packages required for this guide. If you are running this within an internal Google environment where these dependencies are already available, these installation steps can be safely skipped.
Installation#
Install orbax-checkpoint for core checkpointing, flax and optax for the JAX model and optimizer, and orbax-export with tensorflow for exporting to the SavedModel format.
!pip install orbax-checkpoint flax optax
Requirement already satisfied: orbax-checkpoint in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (0.12.5)
Requirement already satisfied: flax in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (0.12.9)
Requirement already satisfied: optax in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (0.2.8)
Requirement already satisfied: absl-py in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (2.5.0)
Requirement already satisfied: etils[epath,epy] in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (1.14.0)
Requirement already satisfied: typing_extensions in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (4.16.0)
Requirement already satisfied: msgpack in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (1.2.2)
Requirement already satisfied: jax>=0.6.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (0.11.2)
Requirement already satisfied: numpy in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (2.5.3)
Requirement already satisfied: prometheus-client>=0.20.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (0.26.0)
Requirement already satisfied: pyyaml in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (6.0.3)
Requirement already satisfied: tensorstore>=0.1.84 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (0.1.85)
Requirement already satisfied: aiofiles in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (25.1.0)
Requirement already satisfied: protobuf in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (7.36.2)
Requirement already satisfied: humanize in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (4.16.0)
Requirement already satisfied: simplejson>=3.16.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (4.1.2)
Requirement already satisfied: psutil in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (7.2.2)
Requirement already satisfied: uvloop in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint) (0.22.1)
Requirement already satisfied: rich>=11.1 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from flax) (15.0.0)
Requirement already satisfied: treescope>=0.1.7 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from flax) (0.1.10)
Requirement already satisfied: jaxlib>=0.5.3 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from optax) (0.11.2)
Requirement already satisfied: ml_dtypes>=0.5.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from jax>=0.6.0->orbax-checkpoint) (0.6.0)
Requirement already satisfied: opt_einsum in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from jax>=0.6.0->orbax-checkpoint) (3.4.0)
Requirement already satisfied: scipy>=1.15 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from jax>=0.6.0->orbax-checkpoint) (1.18.1)
Requirement already satisfied: markdown-it-py>=2.2.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from rich>=11.1->flax) (3.0.0)
Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from rich>=11.1->flax) (2.21.0)
Requirement already satisfied: mdurl~=0.1 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from markdown-it-py>=2.2.0->rich>=11.1->flax) (0.1.2)
Requirement already satisfied: fsspec in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from etils[epath,epy]->orbax-checkpoint) (2026.9.0)
Requirement already satisfied: zipp in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from etils[epath,epy]->orbax-checkpoint) (4.1.0)
!pip install orbax-export tensorflow
Requirement already satisfied: orbax-export in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (0.0.8)
Requirement already satisfied: tensorflow in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (2.20.0rc0)
Requirement already satisfied: absl-py in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-export) (2.5.0)
Requirement already satisfied: dataclasses-json in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-export) (0.6.7)
Requirement already satisfied: etils in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-export) (1.14.0)
Requirement already satisfied: jax>=0.4.34 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-export) (0.11.2)
Requirement already satisfied: jaxlib in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-export) (0.11.2)
Requirement already satisfied: jaxtyping in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-export) (0.3.11)
Requirement already satisfied: numpy in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-export) (2.5.3)
Requirement already satisfied: protobuf in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-export) (7.36.2)
Requirement already satisfied: orbax-checkpoint>=0.9.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-export) (0.12.5)
Requirement already satisfied: astunparse>=1.6.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (1.6.3)
Requirement already satisfied: flatbuffers>=24.3.25 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (25.12.19)
Requirement already satisfied: gast!=0.5.0,!=0.5.1,!=0.5.2,>=0.2.1 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (0.7.0)
Requirement already satisfied: google_pasta>=0.1.1 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (0.2.0)
Requirement already satisfied: libclang>=13.0.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (18.1.1)
Requirement already satisfied: opt_einsum>=2.3.2 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (3.4.0)
Requirement already satisfied: packaging in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (26.3)
Requirement already satisfied: requests<3,>=2.21.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (2.34.2)
Requirement already satisfied: setuptools in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (84.0.0)
Requirement already satisfied: six>=1.12.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (1.17.0)
Requirement already satisfied: termcolor>=1.1.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (3.3.0)
Requirement already satisfied: typing_extensions>=3.6.6 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (4.16.0)
Requirement already satisfied: wrapt>=1.11.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (2.4.1)
Requirement already satisfied: grpcio<2.0,>=1.24.3 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (1.84.0)
Requirement already satisfied: tensorboard~=2.20.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (2.20.0)
Requirement already satisfied: keras>=3.10.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (3.15.1)
Requirement already satisfied: h5py>=3.11.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (3.16.0)
Requirement already satisfied: ml_dtypes<1.0.0,>=0.5.1 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorflow) (0.6.0)
Requirement already satisfied: charset_normalizer<4,>=2 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from requests<3,>=2.21.0->tensorflow) (3.5.1)
Requirement already satisfied: idna<4,>=2.5 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from requests<3,>=2.21.0->tensorflow) (3.20)
Requirement already satisfied: urllib3<3,>=1.26 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from requests<3,>=2.21.0->tensorflow) (2.8.0)
Requirement already satisfied: certifi>=2023.5.7 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from requests<3,>=2.21.0->tensorflow) (2026.7.22)
Requirement already satisfied: markdown>=2.6.8 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorboard~=2.20.0->tensorflow) (3.10.3)
Requirement already satisfied: pillow in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorboard~=2.20.0->tensorflow) (12.3.0)
Requirement already satisfied: tensorboard-data-server<0.8.0,>=0.7.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorboard~=2.20.0->tensorflow) (0.7.2)
Requirement already satisfied: werkzeug>=1.0.1 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from tensorboard~=2.20.0->tensorflow) (3.1.8)
Requirement already satisfied: wheel<1.0,>=0.23.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from astunparse>=1.6.0->tensorflow) (0.40.0)
Requirement already satisfied: scipy>=1.15 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from jax>=0.4.34->orbax-export) (1.18.1)
Requirement already satisfied: rich in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from keras>=3.10.0->tensorflow) (15.0.0)
Requirement already satisfied: namex in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from keras>=3.10.0->tensorflow) (0.1.0)
Requirement already satisfied: optree in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from keras>=3.10.0->tensorflow) (0.20.0)
Requirement already satisfied: msgpack in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint>=0.9.0->orbax-export) (1.2.2)
Requirement already satisfied: prometheus-client>=0.20.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint>=0.9.0->orbax-export) (0.26.0)
Requirement already satisfied: pyyaml in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint>=0.9.0->orbax-export) (6.0.3)
Requirement already satisfied: tensorstore>=0.1.84 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint>=0.9.0->orbax-export) (0.1.85)
Requirement already satisfied: aiofiles in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint>=0.9.0->orbax-export) (25.1.0)
Requirement already satisfied: humanize in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint>=0.9.0->orbax-export) (4.16.0)
Requirement already satisfied: simplejson>=3.16.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint>=0.9.0->orbax-export) (4.1.2)
Requirement already satisfied: psutil in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint>=0.9.0->orbax-export) (7.2.2)
Requirement already satisfied: uvloop in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from orbax-checkpoint>=0.9.0->orbax-export) (0.22.1)
Requirement already satisfied: markupsafe>=2.1.1 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from werkzeug>=1.0.1->tensorboard~=2.20.0->tensorflow) (3.0.3)
Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from dataclasses-json->orbax-export) (3.26.2)
Requirement already satisfied: typing-inspect<1,>=0.4.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from dataclasses-json->orbax-export) (0.9.0)
Requirement already satisfied: mypy-extensions>=0.3.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from typing-inspect<1,>=0.4.0->dataclasses-json->orbax-export) (1.1.0)
Requirement already satisfied: fsspec in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from etils[epath,epy]->orbax-checkpoint>=0.9.0->orbax-export) (2026.9.0)
Requirement already satisfied: zipp in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from etils[epath,epy]->orbax-checkpoint>=0.9.0->orbax-export) (4.1.0)
Requirement already satisfied: wadler-lindig>=0.1.3 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from jaxtyping->orbax-export) (0.1.7)
Requirement already satisfied: markdown-it-py>=2.2.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from rich->keras>=3.10.0->tensorflow) (3.0.0)
Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from rich->keras>=3.10.0->tensorflow) (2.21.0)
Requirement already satisfied: mdurl~=0.1 in /home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages (from markdown-it-py>=2.2.0->rich->keras>=3.10.0->tensorflow) (0.1.2)
Imports#
from orbax.checkpoint import v1 as ocp
import jax
import jax.numpy as jnp
import numpy as np
import flax.linen as nn
import optax
import os
import shutil
from etils import epath
from jax import tree_util
Helper for Directory Management#
A utility function to ensure a clean state for our checkpointing directories during each run of this tutorial.
def cleanup_directory_if_exists(path_str):
"""Removes a directory if it exists."""
path = epath.Path(path_str)
if path.exists():
shutil.rmtree(path)
tutorial_base_dir = epath.Path('/tmp/orbax_tutorial')
cleanup_directory_if_exists(str(tutorial_base_dir))
tutorial_base_dir.mkdir(parents=True, exist_ok=True)
print(f"Tutorial artifacts will be saved under: {tutorial_base_dir}")
Tutorial artifacts will be saved under: /tmp/orbax_tutorial
2. Define a Simulated JAX State#
We’ll construct a PyTree representing our model’s training state. This typically includes model parameters, optimizer state, and the current training step.
Define a Model and Training State#
We will define a basic Flax model, initialize its parameters, and create an Optax optimizer. The complete training state (model parameters, optimizer state, and step count) is stored in a Python dictionary. Sharding is applied to array elements using jax.device_put.
# Model Hyperparameters
input_dim = 64
hidden_dim = 32
output_dim = 10
batch_size_for_init = 4
class SimpleFlaxModel(nn.Module):
hidden_dim: int
output_dim: int
@nn.compact
def __call__(self, x):
x = nn.Dense(features=self.hidden_dim, name="d1")(x)
x = nn.relu(x)
return nn.Dense(features=self.output_dim, name="d2")(x)
key = jax.random.PRNGKey(0)
model_instance = SimpleFlaxModel(hidden_dim, output_dim)
# Initialize model parameters with dummy data.
dummy_input_for_flax_init = jnp.ones((batch_size_for_init, input_dim))
initial_model_params_template = model_instance.init(key, dummy_input_for_flax_init)['params']
np_params = jax.tree_util.tree_map(np.array, initial_model_params_template)
# Initialize the optimizer state.
optimizer_instance = optax.adam(1e-3)
np_opt_state_template = optimizer_instance.init(initial_model_params_template)
# Convert all array-like elements to NumPy arrays, leaving others (like `count`) as-is.
np_opt_state = jax.tree_util.tree_map(lambda x: np.array(x) if hasattr(x, 'shape') else x, np_opt_state_template)
# Define sharding for the model (replicated across all devices).
mesh = jax.sharding.Mesh(jax.devices(), ('data',))
replicated_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
# Group components and apply sharding to all NumPy arrays in the PyTree.
pytree_components_np = {
'params': np_params,
'opt_state': np_opt_state,
}
pytree_components_jax = jax.tree_util.tree_map(lambda x: jax.device_put(x, replicated_sharding) if isinstance(x, np.ndarray) else x, pytree_components_np)
# Combine everything into the final training state PyTree.
simulated_train_state = {**pytree_components_jax, 'step': 0}
print("Initialized JAX training state PyTree with explicit sharding.")
Initialized JAX training state PyTree with explicit sharding.
3. Orbax Checkpointing Workflow#
This section covers managing checkpoints during a simulated training loop using Checkpointer. This API is designed for common training scenarios and allows for powerful configuration through save policies. For a more comprehensive introduction to the Checkpointer API, refer to the Orbax Checkpoint 101 guide.
Create a Checkpointing Directory#
We’ll create a dedicated directory to store our training checkpoints and define a constant for our save interval.
training_ckpt_dir = tutorial_base_dir / 'simulated_training_ckpts'
cleanup_directory_if_exists(str(training_ckpt_dir))
training_ckpt_dir.mkdir(parents=True, exist_ok=True)
SAVE_INTERVAL_STEPS = 2
Checkpointing During a Simulated Training Loop#
We use Checkpointer as a context manager and configure it with a FixedIntervalPolicy. Inside the loop, save(...) is called on every step, but the policy ensures that a checkpoint is only written to disk when the condition (e.g., step % 2 == 0) is met.
# A simplified function to simulate a single training step.
def train_step_for_loop(state):
new_state = state.copy() # Work with a mutable copy of the state dict.
new_state['step'] += 1
# For this demo, we simulate param changes by adding small random noise.
key_for_noise = jax.random.PRNGKey(state['step'])
new_state['params'] = jax.tree_util.tree_map(
lambda p: p + 0.001 * jax.random.normal(key_for_noise, p.shape, p.dtype),
state['params']
)
return new_state
current_loop_state = tree_util.tree_map(lambda x: x, simulated_train_state) # Start with a fresh copy.
num_training_steps = 7
print(f"Simulating {num_training_steps} training steps...")
with ocp.training.Checkpointer(
directory=str(training_ckpt_dir),
save_decision_policy=ocp.training.save_decision_policies.FixedIntervalPolicy(SAVE_INTERVAL_STEPS)
) as ckptr:
for _ in range(num_training_steps):
step_to_save_at = current_loop_state['step']
# `save` takes the current step, the state to save, and optional metrics.
saved = ckptr.save(step_to_save_at, current_loop_state, metrics={'accuracy': 0.85})
if saved: # Will be True if the save_decision_policy decided to save.
print(f" Saved checkpoint for step {step_to_save_at}...")
current_loop_state = train_step_for_loop(current_loop_state)
Simulating 7 training steps...
Saved checkpoint for step 0...
Saved checkpoint for step 2...
Saved checkpoint for step 4...
Saved checkpoint for step 6...
Resuming from a Checkpoint#
To resume training, we use training.Checkpointer.load. Orbax-checkpoint can automatically find the latest completed checkpoint. We provide an abstract_pytree (an empty or example version of our state) to guide the restoration process and ensure the data is loaded with the correct structure and sharding.
with ocp.training.Checkpointer(directory=str(training_ckpt_dir)) as ckptr:
print(f"Restore from the latest checkpoint in {training_ckpt_dir}...")
# It returns None if no checkpoint is found.
resumed_train_state = ckptr.load(
abstract_state=simulated_train_state # Provide an abstract state for structure and sharding.
)
# If a checkpoint was successfully loaded, resumed_train_state will not be None.
if resumed_train_state is not None:
print(f"Restored state successfully. Resuming from step: {resumed_train_state['step']}")
with ocp.training.Checkpointer(directory=str(training_ckpt_dir)) as ckptr:
assert resumed_train_state['step'] == ckptr.latest.step
else:
# If no checkpoint was found, fall back to the initial state.
print("No checkpoint found to restore; using initial state.")
resumed_train_state = simulated_train_state
Restore from the latest checkpoint in /tmp/orbax_tutorial/simulated_training_ckpts...
Restored state successfully. Resuming from step: 6
4. Saving Final JAX Parameters for Export#
After training, you often need to save just the final model parameters for inference or export. For this, Orbax provides the simple save function, which is ideal for one-off saves without the overhead of training policies. See the Checkpointing PyTrees guide for more details on this lower-level API.
Extract Final Parameters for Saving#
We extract the learned parameters from our final training state, as this is the only part we need for inference.
final_params_save_dir = tutorial_base_dir / 'exported_model_params_orbax'
final_model_params_to_save = current_loop_state['params']
print("Final model parameters extracted for saving.")
Final model parameters extracted for saving.
Using save for the Final Save#
save directly saves the given PyTree to the specified directory. It’s a straightforward way to persist the final artifacts of a training process.
# Ensure a clean state by removing the directory if it exists from a previous run.
cleanup_directory_if_exists(str(final_params_save_dir))
print(f"Saving final parameters to: {final_params_save_dir}...")
ocp.save(
path=final_params_save_dir,
state=final_model_params_to_save,
overwrite=True # overwrites an existing checkpoint in directory
)
print("Final model parameters saved via `save`.")
Saving final parameters to: /tmp/orbax_tutorial/exported_model_params_orbax...
Final model parameters saved via `save`.
Loading Exported Parameters (Verification)#
We can use load to load the parameters back and verify that the save operation was successful. Again, we can pass an abstract_pytree to help guide the restoration.
if final_params_save_dir.exists() and len(os.listdir(str(final_params_save_dir))) > 0:
print(f"Loading parameters from {final_params_save_dir} for verification...")
loaded_final_params = ocp.load(
final_params_save_dir,
abstract_state=final_model_params_to_save # Use instance as a template for structure and sharding.
)
# Check that the loaded parameters match the original ones.
params_match = jax.tree_util.tree_all(
jax.tree_util.tree_map(jnp.array_equal, final_model_params_to_save, loaded_final_params)
)
print(f"Verification: {'PASSED' if params_match else 'FAILED'}")
else:
print("Saved parameters directory not found or empty. Skipping verification.")
Loading parameters from /tmp/orbax_tutorial/exported_model_params_orbax for verification...
Verification: PASSED
5. Exporting to TensorFlow SavedModel#
This section demonstrates converting the saved JAX model parameters into a TensorFlow SavedModel format using the orbax export library. This is a common step for for exporting JAX models to TensorFlow SavedModel format.
from orbax.export import ExportManager, JaxModule, ServingConfig
from orbax.export.validate.validation_manager import ValidationManager
import tensorflow as tf
import traceback
import sys
Define JAX Model Apply Function and Pre/Post-processing for Export#
For orbax export, we need to provide a JAX function that takes (params, inputs). We can also define TensorFlow-based pre-processing and post-processing functions, which will be included in the SavedModel’s computation graph.
# `model_instance` was defined in Section 2 (the SimpleFlaxModel instance).
# `final_model_params_to_save` contains the parameters we want to export from Section 4.
# JAX Apply Function: The core JAX logic for the model's forward pass.
@jax.jit
def jax_model_apply_fn_for_export(params, inputs):
"""A JAX function with the signature (params, inputs) for orbax-export."""
return model_instance.apply({'params': params}, inputs)
# Optional: TF Pre-processing Function.
def tf_preprocess_fn_for_export(input_tensor: tf.Tensor) -> tf.Tensor:
"""Normalizes the raw input tensor. Orbax-export will trace this into a graph."""
return tf.cast(input_tensor, tf.float32) / 255.0
# Optional: TF Post-processing Function.
def tf_postprocess_fn_for_export(output_tensor: tf.Tensor) -> dict[str, tf.Tensor]:
"""Packages the model output into a dictionary. Orbax-export will trace this."""
return {'predictions': output_tensor}
print("JAX apply function and plain TF pre/post-processing functions defined for export.")
JAX apply function and plain TF pre/post-processing functions defined for export.
Create JaxModule and ServingConfig#
JaxModule wraps the JAX function and its parameters. ServingConfig defines the input signature for the SavedModel and specifies which pre/post-processing functions to use for a given serving signature key (e.g., serving_default).
# Create the JaxModule, which encapsulates the JAX function and its parameters.
jax_module_for_export = JaxModule(
params=final_model_params_to_save,
apply_fn=jax_model_apply_fn_for_export,
input_polymorphic_shape=f'(b, {input_dim})',
jax2tf_kwargs={'with_gradient': False, 'native_serialization': False}
)
# This tells orbax-export how to trace the Python preprocessor function.
tf_input_signature = [
tf.TensorSpec(shape=[None, input_dim], dtype=tf.float32)
]
# Create a serving configuration that bundles the signature key, input specs,
# and our Python processing functions.
serving_config = ServingConfig(
signature_key='serving_default',
input_signature=tf_input_signature,
tf_preprocessor=tf_preprocess_fn_for_export,
tf_postprocessor=tf_postprocess_fn_for_export
)
print("JaxModule and ServingConfig created successfully.")
2026-09-22 17:08:02.734091: E external/local_xla/xla/stream_executor/cuda/cuda_platform.cc:51] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: UNKNOWN ERROR (303)
JaxModule and ServingConfig created successfully.
Export to TensorFlow SavedModel#
The ExportManager takes the JaxModule and a list of ServingConfig to build and save the final TensorFlow SavedModel.
# Define the directory to save the final exported model.
saved_model_dir = tutorial_base_dir / 'tf_saved_model_orbax_export'
cleanup_directory_if_exists(str(saved_model_dir))
saved_model_dir.mkdir(parents=True, exist_ok=True)
# The ExportManager orchestrates the JAX-to-TF conversion and saving process.
export_manager = ExportManager(jax_module_for_export, [serving_config])
print(f"Exporting SavedModel to: {saved_model_dir}")
try:
export_manager.save(str(saved_model_dir))
print("Model exported successfully to SavedModel format.")
print(f"Contents of {saved_model_dir}: {os.listdir(str(saved_model_dir))}")
except Exception as e:
print(f"ERROR during SavedModel export: {e}")
import traceback
traceback.print_exc()
Exporting SavedModel to: /tmp/orbax_tutorial/tf_saved_model_orbax_export
Model exported successfully to SavedModel format.
Contents of /tmp/orbax_tutorial/tf_saved_model_orbax_export: ['fingerprint.pb', 'assets', 'saved_model.pb', 'variables']
Validate the Exported Model#
A critical final step is to verify that the exported TensorFlow model produces the same results as the original JAX model. We use the ValidationManager, which compares the outputs of the JAX model and the loaded TF SavedModel for a given batch of inputs and generates a detailed report.
# Prepare a batch of test inputs. These should be "raw" (pre-preprocessing).
validation_batch_size = 4
raw_validation_inputs = np.random.rand(validation_batch_size, input_dim).astype(np.float32) * 255.0
# To match the positional signature pass inputs as a list of lists.
validation_mgr = ValidationManager(
module=jax_module_for_export,
serving_configs=[serving_config],
model_inputs=[[raw_validation_inputs]]
)
# Load the candidate model we want to validate.
loaded_tf_model = tf.saved_model.load(str(saved_model_dir))
# Run the validation, which compares the JAX and TF outputs.
print("\nRunning validation...")
validation_reports = validation_mgr.validate(loaded_tf_model)
# Check the report. The report is a dict keyed by the signature_key.
report = validation_reports['serving_default']
# The report status is an enum. We check its string name for a simple pass/fail result.
if report.status.name == 'Pass':
print(f"VERIFICATION PASSED! Status: {report.status.name}")
else:
print(f"VERIFICATION FAILED! Status: {report.status.name}")
# The report can be printed as a JSON string for detailed inspection of differences and latencies.
print("\nValidation Report:")
print(report.to_json(indent=2))
Running validation...
2026-09-22 17:08:03.305474: W tensorflow/core/framework/op_kernel.cc:1831] OP_REQUIRES failed at xla_call_module_op.cc:225 : INVALID_ARGUMENT: Cannot deserialize computation: UNKNOWN: <unknown>:0: error: loc("shape_assertion"): unregistered operation 'vhlo.custom_call_v2' found in dialect ('vhlo') that does not allow unknown operations
<unknown>:0: note: loc("shape_assertion"): see current operation: "vhlo.custom_call_v2"(%23, %22) <#vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>> {error_message = #vhlo.string_v1<"Input shapes do not match the polymorphic shapes specification. Expected value >= 1 for dimension variable 'b'. Using the following polymorphic shapes specifications: args[1].shape = (b, 64). Obtained dimension variables: 'b' = {0} from specification 'b' for dimension args[1].shape[0] (= {0}), . Please see https://docs.jax.dev/en/latest/export/shape_poly.html#shape-assertion-errors for more details.">} : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> ()
<unknown>:0: note: loc("shape_assertion"): in bytecode version 6 produced by: StableHLO_v1.18.0
<unknown>:0: error: failed to deserialize portable artifact using StableHLO_v1.12.1
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1790096883.305716 3168 graph_compiler.cc:153] Executor failed to create kernel. INVALID_ARGUMENT: Cannot deserialize computation: UNKNOWN: <unknown>:0: error: loc("shape_assertion"): unregistered operation 'vhlo.custom_call_v2' found in dialect ('vhlo') that does not allow unknown operations
<unknown>:0: note: loc("shape_assertion"): see current operation: "vhlo.custom_call_v2"(%23, %22) <#vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>> {error_message = #vhlo.string_v1<"Input shapes do not match the polymorphic shapes specification. Expected value >= 1 for dimension variable 'b'. Using the following polymorphic shapes specifications: args[1].shape = (b, 64). Obtained dimension variables: 'b' = {0} from specification 'b' for dimension args[1].shape[0] (= {0}), . Please see https://docs.jax.dev/en/latest/export/shape_poly.html#shape-assertion-errors for more details.">} : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> ()
<unknown>:0: note: loc("shape_assertion"): in bytecode version 6 produced by: StableHLO_v1.18.0
<unknown>:0: error: failed to deserialize portable artifact using StableHLO_v1.12.1
[[{{node XlaCallModule}}]]
2026-09-22 17:08:03.305806: W tensorflow/core/framework/op_kernel.cc:1855] OP_REQUIRES failed at xla_ops.cc:590 : INVALID_ARGUMENT: Cannot deserialize computation: UNKNOWN: <unknown>:0: error: loc("shape_assertion"): unregistered operation 'vhlo.custom_call_v2' found in dialect ('vhlo') that does not allow unknown operations
<unknown>:0: note: loc("shape_assertion"): see current operation: "vhlo.custom_call_v2"(%23, %22) <#vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>> {error_message = #vhlo.string_v1<"Input shapes do not match the polymorphic shapes specification. Expected value >= 1 for dimension variable 'b'. Using the following polymorphic shapes specifications: args[1].shape = (b, 64). Obtained dimension variables: 'b' = {0} from specification 'b' for dimension args[1].shape[0] (= {0}), . Please see https://docs.jax.dev/en/latest/export/shape_poly.html#shape-assertion-errors for more details.">} : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> ()
<unknown>:0: note: loc("shape_assertion"): in bytecode version 6 produced by: StableHLO_v1.18.0
<unknown>:0: error: failed to deserialize portable artifact using StableHLO_v1.12.1
[[{{node XlaCallModule}}]]
tf2xla conversion failed while converting __inference_predict_fn_191[]. Run with TF_DUMP_GRAPH_PREFIX=/path/to/dump/dir and --vmodule=xla_compiler=2 to obtain a dump of the compiled functions.
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/function_type_utils.py:442, in bind_function_inputs(args, kwargs, function_type, default_values)
441 try:
--> 442 bound_arguments = function_type.bind_with_defaults(
443 args, sanitized_kwargs, default_values
444 )
445 except Exception as e:
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/core/function/polymorphism/function_type.py:264, in FunctionType.bind_with_defaults(self, args, kwargs, default_values)
263 """Returns BoundArguments with default values filled in."""
--> 264 bound_arguments = self.bind(*args, **kwargs)
265 bound_arguments.apply_defaults()
File ~/.asdf/installs/python/3.12.13/lib/python3.12/inspect.py:3280, in Signature.bind(self, *args, **kwargs)
3276 """Get a BoundArguments object, that maps the passed `args`
3277 and `kwargs` to the function's signature. Raises `TypeError`
3278 if the passed arguments can not be bound.
3279 """
-> 3280 return self._bind(args, kwargs)
File ~/.asdf/installs/python/3.12.13/lib/python3.12/inspect.py:3204, in Signature._bind(self, args, kwargs, partial)
3201 if param.kind in (_VAR_KEYWORD, _KEYWORD_ONLY):
3202 # Looks like we have no parameter for this positional
3203 # argument
-> 3204 raise TypeError(
3205 'too many positional arguments') from None
3207 if param.kind == _VAR_POSITIONAL:
3208 # We have an '*args'-like argument, let's fill it with
3209 # all positional arguments we have left and move on to
3210 # the next phase
TypeError: too many positional arguments
The above exception was the direct cause of the following exception:
TypeError Traceback (most recent call last)
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/concrete_function.py:1179, in ConcreteFunction._call_impl(self, args, kwargs)
1178 try:
-> 1179 return self._call_with_structured_signature(args, kwargs)
1180 except TypeError as structured_err:
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/concrete_function.py:1259, in ConcreteFunction._call_with_structured_signature(self, args, kwargs)
1245 """Executes the wrapped function with the structured signature.
1246
1247 Args:
(...) 1256 of this `ConcreteFunction`.
1257 """
1258 bound_args = (
-> 1259 function_type_utils.canonicalize_function_inputs(
1260 args, kwargs, self.function_type)
1261 )
1262 filtered_flat_args = self.function_type.unpack_inputs(bound_args)
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/function_type_utils.py:422, in canonicalize_function_inputs(args, kwargs, function_type, default_values, is_pure)
421 args, kwargs = _convert_variables_to_tensors(args, kwargs)
--> 422 bound_arguments = bind_function_inputs(
423 args, kwargs, function_type, default_values
424 )
425 return bound_arguments
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/function_type_utils.py:446, in bind_function_inputs(args, kwargs, function_type, default_values)
445 except Exception as e:
--> 446 raise TypeError(
447 f"Binding inputs to tf.function failed due to `{e}`. "
448 f"Received args: {args} and kwargs: {sanitized_kwargs} for signature:"
449 f" {function_type}."
450 ) from e
451 return bound_arguments
TypeError: Binding inputs to tf.function failed due to `too many positional arguments`. Received args: (<tf.Tensor: shape=(4, 64), dtype=float32, numpy=
array([[154.91234 , 197.58281 , 127.42949 , 108.02718 ,
90.35517 , 157.41173 , 185.41382 , 252.16258 ,
6.7797484 , 213.46227 , 197.31313 , 129.75708 ,
46.119354 , 136.47139 , 138.64201 , 22.383392 ,
97.96809 , 180.20482 , 79.75496 , 179.16927 ,
140.3409 , 92.693405 , 110.860565 , 67.73705 ,
220.61707 , 215.28065 , 12.4419775 , 95.83215 ,
88.08871 , 225.45192 , 124.54951 , 133.94673 ,
131.8363 , 160.95447 , 240.34119 , 66.9754 ,
46.000183 , 15.57432 , 228.45929 , 194.90842 ,
163.85056 , 114.81381 , 66.98274 , 122.93302 ,
48.449455 , 216.43806 , 165.15515 , 3.7304585 ,
197.51535 , 128.59877 , 190.22537 , 246.28343 ,
95.28822 , 42.412605 , 31.931856 , 216.43393 ,
233.02641 , 95.12167 , 205.99136 , 233.52435 ,
64.42335 , 203.00887 , 117.244705 , 54.250954 ],
[ 80.725266 , 135.81294 , 61.912308 , 64.61787 ,
211.56705 , 87.94533 , 70.99094 , 10.107669 ,
11.231559 , 145.35521 , 160.23227 , 16.030617 ,
76.98017 , 52.703064 , 197.10526 , 114.18886 ,
4.8712273 , 133.80534 , 200.19373 , 106.54956 ,
220.74739 , 31.23908 , 59.54667 , 212.6548 ,
41.286728 , 35.65776 , 30.72895 , 30.615042 ,
109.18758 , 168.65205 , 116.9906 , 155.1923 ,
45.11585 , 254.59741 , 13.33656 , 4.3989053 ,
245.3573 , 29.116545 , 213.34816 , 164.61497 ,
93.96884 , 187.84195 , 225.26126 , 0.27181178,
79.4282 , 223.67604 , 27.164669 , 229.07487 ,
179.93153 , 145.62474 , 236.37834 , 84.19193 ,
219.93732 , 218.375 , 230.24217 , 72.67611 ,
56.489452 , 87.918274 , 173.30745 , 214.71097 ,
59.796444 , 120.88078 , 209.4332 , 137.33475 ],
[191.43623 , 51.953697 , 123.888885 , 11.992398 ,
101.470116 , 111.174736 , 194.57674 , 212.18513 ,
132.04231 , 120.8159 , 126.571434 , 82.43639 ,
141.6653 , 162.5258 , 229.3476 , 63.817863 ,
161.2689 , 247.80777 , 38.901596 , 81.91015 ,
101.064545 , 213.38788 , 94.15804 , 117.10467 ,
139.95174 , 110.55954 , 28.893095 , 91.79001 ,
10.967405 , 13.654747 , 131.08783 , 67.01941 ,
166.53772 , 19.98414 , 235.02274 , 4.201903 ,
145.8828 , 141.33704 , 154.95238 , 38.97124 ,
231.74254 , 176.93297 , 74.18795 , 171.45912 ,
188.7082 , 243.04433 , 184.41545 , 60.996906 ,
249.6386 , 47.6212 , 193.00038 , 232.88686 ,
224.61368 , 151.62221 , 219.92427 , 127.420166 ,
162.34135 , 229.24112 , 107.78789 , 194.26643 ,
28.373556 , 0.37842697, 64.019104 , 93.44569 ],
[175.78894 , 81.426315 , 49.171825 , 1.4068505 ,
117.60167 , 98.26583 , 69.77691 , 96.12167 ,
106.3485 , 121.31828 , 110.58626 , 251.60945 ,
123.64582 , 166.34799 , 170.22653 , 83.12016 ,
184.84718 , 209.5949 , 144.0851 , 55.771355 ,
68.847984 , 224.4301 , 158.94998 , 184.65622 ,
49.27672 , 59.845867 , 51.994038 , 169.9199 ,
193.33105 , 179.18019 , 128.99184 , 249.57642 ,
212.21123 , 177.64696 , 34.46966 , 126.478065 ,
142.8804 , 164.4711 , 54.35215 , 105.15004 ,
53.885105 , 170.48935 , 236.80267 , 113.81077 ,
136.55412 , 72.96272 , 244.60863 , 12.582669 ,
99.78536 , 63.913902 , 194.64223 , 46.225163 ,
10.036768 , 114.148994 , 12.067425 , 104.96419 ,
175.71404 , 57.019577 , 118.73363 , 0.3565527 ,
231.80464 , 55.94336 , 62.73303 , 100.582085 ]],
dtype=float32)>,) and kwargs: {} for signature: (*, inputs_0: TensorSpec(shape=(None, 64), dtype=tf.float32, name='inputs_0')) -> Dict[['predictions', TensorSpec(shape=(None, 10), dtype=tf.float32, name='predictions')]].
During handling of the above exception, another exception occurred:
InvalidArgumentError Traceback (most recent call last)
Cell In[16], line 17
15 # Run the validation, which compares the JAX and TF outputs.
16 print("\nRunning validation...")
---> 17 validation_reports = validation_mgr.validate(loaded_tf_model)
19 # Check the report. The report is a dict keyed by the signature_key.
20 report = validation_reports['serving_default']
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/orbax/export/validate/validation_manager.py:201, in ValidationManager.validate(self, loaded_model, with_xprof, report_option)
197 validation_job = ValidationJob(
198 baseline_fns[key], candidate_fns[key], input_map[key], with_xprof
199 )
200 baseline_result = validation_job.calc_baseline_result()
--> 201 candidate_result = validation_job.calc_candidate_result()
202 # Always convert list to Dict
203 baseline_result.maybe_convert_result_to_dict()
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/orbax/export/validate/validation_job.py:126, in ValidationJob.calc_candidate_result(self)
124 def calc_candidate_result(self) -> ValidationSingleJobResult:
125 """Feed batch_input into candidate `inference_fn` and run."""
--> 126 return self._calc_result(self._candidate_inference_fn, self._batch_input)
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/orbax/export/validate/validation_job.py:98, in ValidationJob._calc_result(self, infer_step, batched_examples)
96 """Feed batch_input into `apply_fn` and run."""
97 # Warm up in case they are jit functions.
---> 98 _ = infer_step(batched_examples[0])
100 # Generate baseline model baseline result.
101 latencies = list()
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/orbax/export/validate/validation_manager.py:141, in ValidationManager._create_candidate_fns.<locals>.make_candidate_inference_fn.<locals>.inference_fn(*inputs)
139 outputs = loaded_model_signatures[signature_key](**real_inputs)
140 elif isinstance(real_inputs, Sequence):
--> 141 outputs = loaded_model_signatures[signature_key](*real_inputs)
142 else:
143 outputs = loaded_model_signatures[signature_key](real_inputs)
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/concrete_function.py:1170, in ConcreteFunction.__call__(self, *args, **kwargs)
1120 def __call__(self, *args, **kwargs):
1121 """Executes the wrapped function.
1122
1123 ConcreteFunctions have two signatures:
(...) 1168 TypeError: If the arguments do not match the function's signature.
1169 """
-> 1170 return self._call_impl(args, kwargs)
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/concrete_function.py:1182, in ConcreteFunction._call_impl(self, args, kwargs)
1180 except TypeError as structured_err:
1181 try:
-> 1182 return self._call_with_flat_signature(args, kwargs)
1183 except (TypeError, ValueError) as flat_err:
1184 raise TypeError( # pylint: disable=raise-missing-from
1185 str(structured_err)
1186 + "\nFallback to flat signature also failed due to: "
1187 + str(flat_err)
1188 )
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/concrete_function.py:1242, in ConcreteFunction._call_with_flat_signature(self, args, kwargs)
1237 if not isinstance(
1238 arg, (tensor_lib.Tensor, resource_variable_ops.BaseResourceVariable)):
1239 raise TypeError(f"{self._flat_signature_summary()}: expected argument "
1240 f"#{i}(zero-based) to be a Tensor; "
1241 f"got {type(arg).__name__} ({arg}).")
-> 1242 return self._call_flat(args, self.captured_inputs)
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/saved_model/load.py:146, in _WrapperFunction._call_flat(self, args, captured_inputs)
144 else: # cross-replica context
145 captured_inputs = list(map(get_unused_handle, captured_inputs))
--> 146 return super()._call_flat(args, captured_inputs)
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/concrete_function.py:1322, in ConcreteFunction._call_flat(self, tensor_inputs, captured_inputs)
1318 possible_gradient_type = gradients_util.PossibleTapeGradientTypes(args)
1319 if (possible_gradient_type == gradients_util.POSSIBLE_GRADIENT_TYPES_NONE
1320 and executing_eagerly):
1321 # No tape is watching; skip to running the function.
-> 1322 return self._inference_function.call_preflattened(args)
1323 forward_backward = self._select_forward_and_backward_functions(
1324 args,
1325 possible_gradient_type,
1326 executing_eagerly)
1327 forward_function, args_with_tangents = forward_backward.forward()
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/atomic_function.py:216, in AtomicFunction.call_preflattened(self, args)
214 def call_preflattened(self, args: Sequence[core.Tensor]) -> Any:
215 """Calls with flattened tensor inputs and returns the structured output."""
--> 216 flat_outputs = self.call_flat(*args)
217 return self.function_type.pack_output(flat_outputs)
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/atomic_function.py:251, in AtomicFunction.call_flat(self, *args)
249 with record.stop_recording():
250 if self._bound_context.executing_eagerly():
--> 251 outputs = self._bound_context.call_function(
252 self.name,
253 list(args),
254 len(self.function_type.flat_outputs),
255 )
256 else:
257 outputs = make_call_op_in_graph(
258 self,
259 list(args),
260 self._bound_context.function_call_options.as_attrs(),
261 )
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/context.py:1688, in Context.call_function(self, name, tensor_inputs, num_outputs)
1686 cancellation_context = cancellation.context()
1687 if cancellation_context is None:
-> 1688 outputs = execute.execute(
1689 name.decode("utf-8"),
1690 num_outputs=num_outputs,
1691 inputs=tensor_inputs,
1692 attrs=attrs,
1693 ctx=self,
1694 )
1695 else:
1696 outputs = execute.execute_with_cancellation(
1697 name.decode("utf-8"),
1698 num_outputs=num_outputs,
(...) 1702 cancellation_manager=cancellation_context,
1703 )
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/tensorflow/python/eager/execute.py:53, in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
51 try:
52 ctx.ensure_initialized()
---> 53 tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
54 inputs, attrs, num_outputs)
55 except core._NotOkStatusException as e:
56 if name is not None:
InvalidArgumentError: Graph execution error:
Detected at node XlaCallModule defined at (most recent call last):
<stack traces unavailable>
Cannot deserialize computation: UNKNOWN: <unknown>:0: error: loc("shape_assertion"): unregistered operation 'vhlo.custom_call_v2' found in dialect ('vhlo') that does not allow unknown operations
<unknown>:0: note: loc("shape_assertion"): see current operation: "vhlo.custom_call_v2"(%23, %22) <#vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>> {error_message = #vhlo.string_v1<"Input shapes do not match the polymorphic shapes specification. Expected value >= 1 for dimension variable 'b'. Using the following polymorphic shapes specifications: args[1].shape = (b, 64). Obtained dimension variables: 'b' = {0} from specification 'b' for dimension args[1].shape[0] (= {0}), . Please see https://docs.jax.dev/en/latest/export/shape_poly.html#shape-assertion-errors for more details.">} : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> ()
<unknown>:0: note: loc("shape_assertion"): in bytecode version 6 produced by: StableHLO_v1.18.0
<unknown>:0: error: failed to deserialize portable artifact using StableHLO_v1.12.1
[[{{node XlaCallModule}}]]
tf2xla conversion failed while converting __inference_predict_fn_191[]. Run with TF_DUMP_GRAPH_PREFIX=/path/to/dump/dir and --vmodule=xla_compiler=2 to obtain a dump of the compiled functions.
[[StatefulPartitionedCall/StatefulPartitionedCall]] [Op:__inference_signature_wrapper_inference_fn_254]