Export and serve a trained JAX model

1. Introduction

Jax on GPU Learning path. Lab 8: Serving and Next Steps.

In this codelab you start from a trained checkpoint and walk the JAX inference pipeline end to end, from a JIT-compiled forward pass to ahead-of-time (AOT) compilation, JAX-native export with jax.export, and a TensorFlow SavedModel with jax2tf. Each path is measured, and all four are meant to produce the same predictions.

What you'll do

Rebuild the codelab 7 transformer and load its trained weights with Orbax
Wrap the forward pass in jax.jit and measure first-call latency against cached-call latency
Remove the cold start with AOT compilation (lower() then compile()) and read the StableHLO IR
Measure forward-pass throughput in tokens/sec across four batch sizes
Serialize a portable artifact with jax.export, then deserialize and call it
Convert the model to a TensorFlow SavedModel with jax2tf and compare all four paths

What you'll need

A Google Cloud project with billing enabled, and workshop credits or a reservation covering GPU usage
Quota for at least 2 NVIDIA L4 GPUs in your chosen region (how to check GPU quota)
Completion of codelabs 1 to 7, or an equivalent CUDA-enabled JAX GPU environment
A completed run of codelab 7, whose Orbax checkpoint this lab loads

Estimated time to complete: 60 minutes.

Four paths from training to serving

This codelab compares four practical ways to move a trained JAX model toward serving, depending on your deployment target:

Path	Format	Serving target	When to use
`jax.jit`	Cached in-process executable	Python server (FastAPI, Flask)	Simplest low-latency JAX serving path
`AOT compile`	Precompiled in-process executable	Python server startup / warmup	Avoid first-request compilation latency
`jax.export`	Serialized JAX export with StableHLO + metadata	Compatible JAX runtime for the exported platform(s)	JAX-native portable artifact
`jax2tf`	TF SavedModel	TF Serving, TFX	TensorFlow ecosystem

You run all four, starting from the same trained checkpoint. The differences are in portability, dependencies, and whether the artifact can be used outside the original Python process.

2. Before you begin

Select your project

In the Google Cloud console, select or create a project with billing enabled.

Open Cloud Shell

Click Activate Cloud Shell (the terminal icon at the top right of the console) to start a Cloud Shell session, then point it at your project:

gcloud config set project <YOUR_PROJECT_ID>

Provision the GPU environment

Run the following in Cloud Shell. Codelab 1 explains each command in detail, including GPU quota and zone requirements.

gcloud services enable \
  container.googleapis.com \
  compute.googleapis.com \
  iam.googleapis.com \
  cloudresourcemanager.googleapis.com \
  logging.googleapis.com \
  monitoring.googleapis.com

git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu/terraform

cp terraform.tfvars.example terraform.tfvars

Edit terraform.tfvars and set project_id = "", then provision the cluster and deploy JupyterLab:

terraform init
terraform apply
$(terraform output -raw get_credentials_command)

cd ..
kubectl apply -f deploy/jupyter.yaml

terraform apply takes about 12 minutes. When it finishes, wait for the Pod and the LoadBalancer, then read the one-time JupyterLab token from the Pod log:

kubectl get pod jax-jupyter -w         # wait for Running, then Ctrl+C
kubectl get svc jax-jupyter-svc -w     # wait for EXTERNAL-IP, then Ctrl+C
kubectl logs jax-jupyter | grep -o 'token=[a-z0-9]*' | head -1

Open http://:8884, paste the token, and create a new Python 3 notebook in /workspace. Every code block in this codelab goes into a cell of that notebook.

Install what this codelab needs

!pip install --quiet flax orbax-checkpoint tensorflow

flax and orbax-checkpoint usually ship in the NVIDIA JAX container, so those two are normally a no-op. If TensorFlow is not already present, pip pulls it.

Set up and verify the GPU

Import JAX, Flax NNX, and Orbax, confirm that at least one GPU is visible, and define the two helpers the rest of the codelab uses for blocking on results and rendering result tables.

import os
os.environ["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib64:" + os.environ.get("LD_LIBRARY_PATH", "")

import html
import pathlib
import time
import warnings

from IPython.display import HTML, display
import numpy as np

warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", message=".*ml_dtypes.*")
warnings.filterwarnings("ignore", message=".*JAX_PLATFORMS.*")

import jax
import jax.numpy as jnp
from flax import nnx
import orbax.checkpoint as ocp


devices = jax.devices()
gpu_devices = [d for d in devices if d.platform == "gpu"]

print(f"JAX version:     {jax.__version__}")
print(f"Default backend: {jax.default_backend()}")
print(f"GPU devices:     {gpu_devices}")

assert len(gpu_devices) >= 1, (
    f"This lesson needs at least 1 GPU. Found {len(gpu_devices)}. "
    f"Available devices: {devices}"
)


def block_tree(tree):
    return jax.block_until_ready(tree)


def show_table(headers, rows, title=None, aligns=None):
    aligns = aligns or ["left"] * len(headers)
    parts = ["<div style='font-family: system-ui; max-width: 980px;'>"]
    if title:
        parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    parts.append("<table style='border-collapse: collapse; width: 100%; font-size: 13px;'>")
    parts.append("<thead><tr>")
    for h, a in zip(headers, aligns):
        parts.append(
            f"<th style='text-align:{a}; border-bottom:1px solid #d0d7de; padding:6px;'>"
            f"{html.escape(str(h))}</th>"
        )
    parts.append("</tr></thead><tbody>")
    for row in rows:
        parts.append("<tr>")
        for cell, a in zip(row, aligns):
            parts.append(
                f"<td style='text-align:{a}; border-bottom:1px solid #eef1f4; padding:6px;'>"
                f"{html.escape(str(cell))}</td>"
            )
        parts.append("</tr>")
    parts.append("</tbody></table></div>")
    display(HTML("\n".join(parts)))

You should see a JAX version, gpu as the default backend, and a list with at least one CUDA device.

Confirm you have a trained checkpoint

This codelab loads the Orbax checkpoint that Train a transformer end to end with Flax NNX and Orbax codelab wrote to /tmp/jax-course/l7-checkpoints/trained, so check that the directory is still there before you build anything.

import pathlib
ckpt_dir = pathlib.Path("/tmp/jax-course/l7-checkpoints")
assert (ckpt_dir / "trained").exists(), (
    f"No checkpoint at {ckpt_dir / 'trained'}. Run codelab 7 first."
)
print(f"Found checkpoint: {ckpt_dir / 'trained'}")

You should see the checkpoint path printed back to you.

3. Rebuild the model and load the checkpoint

An Orbax checkpoint stores parameter values, not the model class that produced them. To restore it you first need the same architecture, so this step re-declares the codelab 7 TinyTransformer and then fills it with the saved weights.

Redefine the architecture

The model has a token and position embedding, four pre-norm transformer blocks with causal attention, a final layer norm, and a language-model head.

VOCAB_SIZE = 256
D_MODEL = 256
NUM_HEADS = 4
FFN_DIM = 1024
NUM_LAYERS = 4
MAX_SEQ_LEN = 256


def causal_sdpa(query, key, value, **_):
    return jax.nn.dot_product_attention(query, key, value, is_causal=True)


class TransformerBlock(nnx.Module):
    def __init__(self, d_model: int, num_heads: int, ffn_dim: int, rngs: nnx.Rngs):
        self.ln1 = nnx.LayerNorm(d_model, rngs=rngs)
        self.attn = nnx.MultiHeadAttention(
            num_heads=num_heads,
            in_features=d_model,
            decode=False,
            attention_fn=causal_sdpa,
            rngs=rngs,
        )
        self.ln2 = nnx.LayerNorm(d_model, rngs=rngs)
        self.fc_up = nnx.Linear(d_model, ffn_dim, rngs=rngs)
        self.fc_down = nnx.Linear(ffn_dim, d_model, rngs=rngs)

    def __call__(self, x):
        x = x + self.attn(self.ln1(x))
        h = jax.nn.gelu(self.fc_up(self.ln2(x)))
        x = x + self.fc_down(h)
        return x


class TinyTransformer(nnx.Module):
    def __init__(self, vocab_size: int, d_model: int, num_heads: int,
                 ffn_dim: int, num_layers: int, max_seq_len: int,
                 rngs: nnx.Rngs):
        self.token_embed = nnx.Embed(vocab_size, d_model, rngs=rngs)
        self.pos_embed = nnx.Embed(max_seq_len, d_model, rngs=rngs)
        self.blocks = nnx.List([
            TransformerBlock(d_model, num_heads, ffn_dim, rngs=rngs)
            for _ in range(num_layers)
        ])
        self.final_norm = nnx.LayerNorm(d_model, rngs=rngs)
        self.lm_head = nnx.Linear(d_model, vocab_size, use_bias=False, rngs=rngs)

    def __call__(self, tokens):
        B, T = tokens.shape
        x = self.token_embed(tokens) + self.pos_embed(jnp.arange(T))
        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        return self.lm_head(x)


param_count = sum(x.size for x in jax.tree.leaves(nnx.state(TinyTransformer(
    VOCAB_SIZE, D_MODEL, NUM_HEADS, FFN_DIM, NUM_LAYERS, MAX_SEQ_LEN,
    rngs=nnx.Rngs(0),
), nnx.Param)))
print(f"TinyTransformer: {param_count:,} parameters")

You should see a parameter count for a freshly initialized model. The weights are random at this point because the architecture is what matters here.

Restore the trained weights

Orbax restores into a target structure. You build that target by taking the state of the fresh model and replacing every array with a jax.ShapeDtypeStruct, which describes a shape and dtype without allocating memory. Orbax then reads the checkpoint into arrays of exactly those shapes.

ckpt_dir = pathlib.Path("/tmp/jax-course/l7-checkpoints")

model = TinyTransformer(
    VOCAB_SIZE, D_MODEL, NUM_HEADS, FFN_DIM, NUM_LAYERS, MAX_SEQ_LEN,
    rngs=nnx.Rngs(0),
)
model_params = nnx.state(model, nnx.Param)
abstract_params = jax.tree.map(
    lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), model_params
)

checkpointer = ocp.StandardCheckpointer()
restored_params = checkpointer.restore(ckpt_dir / "trained", abstract_params)

# Move to a single GPU
single_device = gpu_devices[0]
restored_params = jax.device_put(restored_params, single_device)
nnx.update(model, restored_params)

print(f"\u2705 Checkpoint loaded from {ckpt_dir / 'trained'}")
print(f"Parameters: {sum(x.size for x in jax.tree.leaves(restored_params)):,}")

The jax.device_put call matters. Serving uses one GPU, so the parameters are moved onto gpu_devices[0] before they go into the model.

You should see the checkpoint path confirmed and a parameter count that matches the one printed by the previous cell.

4. Path 1: serve with jax.jit

The simplest serving path is wrapping the forward pass in jax.jit. The first call triggers compilation and every call after that reuses the cached executable.

For serving, bake the weights into a closure. nnx.split separates the model into a static graph definition and dynamic state and the jitted function captures both and takes only tokens as an argument. That makes the compiled function self-contained.

graphdef, model_state = nnx.split(model)

@jax.jit
def predict_jit(tokens):
    m = nnx.merge(graphdef, model_state)
    return m(tokens)

dummy_input = jnp.zeros((1, MAX_SEQ_LEN), dtype=jnp.int32)

start = time.perf_counter()
logits = block_tree(predict_jit(dummy_input))
first_call_ms = (time.perf_counter() - start) * 1000

times = []
for _ in range(100):
    start = time.perf_counter()
    logits = block_tree(predict_jit(dummy_input))
    times.append((time.perf_counter() - start) * 1000)

avg_ms = np.mean(times)

show_table(
    ["", "Latency (ms)"],
    [
        ("First call (compile + execute)", f"{first_call_ms:,.1f}"),
        ("Subsequent calls (avg of 100)", f"{avg_ms:.2f}"),
        ("Speedup", f"{first_call_ms / avg_ms:.0f}\u00d7"),
    ],
    title="JIT inference latency",
    aligns=["left", "right"],
)

print(f"\nOutput shape: {logits.shape} (batch=1, seq={MAX_SEQ_LEN}, vocab={VOCAB_SIZE})")

You should see a JIT inference latency table in which the first call is much slower than the average of the following 100, and then Output shape: (1, 256, 256).

5. Path 2: remove the cold start with AOT compilation

jax.jit compiles on first call, which is fine during development and a problem in serving. AOT compilation splits that single step into separate stages:

You write a normal Python/JAX function, such as a model forward pass
JAX traces the function for specific input shapes and dtypes and converts it into compiler IR using lower()
StableHLO describes the computation in hardware-independent ops like dot, reshape, and reduce.
XLA optimizes the StableHLO and produces a device-specific executable for GPU, TPU, or CPU with compile().

StableHLO intermediate representation (IR) is the compiler-level representation of a JAX computation after Python code has been lowered into a portable, hardware-independent program. It describes operations such as matrix multiplies, reshapes, reductions, and control flow in a form that XLA can compile for different backends, including GPUs, TPUs, and CPUs.

You compile once, either at startup, or offline for a fixed input shape and then execute many times without paying first-call compilation overhead.

# Bake weights into a closure
def predict_closed(tokens):
    m = nnx.merge(graphdef, model_state)
    return m(tokens)

# Stage 1: Lower
abstract_tokens = jax.ShapeDtypeStruct((1, MAX_SEQ_LEN), jnp.int32)

lowered = predict_closed.lower(abstract_tokens)
print(f"Lowered to StableHLO ({len(lowered.as_text()):,} chars)")

# Stage 2: Compile
compiled = lowered.compile()
print(f"Compiled for: {jax.default_backend()}")

# Execute
start = time.perf_counter()
logits_aot = block_tree(compiled(dummy_input))
aot_first_ms = (time.perf_counter() - start) * 1000

times_aot = []
for _ in range(100):
    start = time.perf_counter()
    logits_aot = block_tree(compiled(dummy_input))
    times_aot.append((time.perf_counter() - start) * 1000)

avg_aot_ms = np.mean(times_aot)
max_diff_jit_aot = float(jnp.max(jnp.abs(logits - logits_aot)))

show_table(
    ["", "Latency (ms)"],
    [
        ("AOT first execution (no compile)", f"{aot_first_ms:.2f}"),
        ("AOT subsequent (avg of 100)", f"{avg_aot_ms:.2f}"),
        ("JIT first call (from above)", f"{first_call_ms:,.1f}"),
        ("Max |JIT − AOT|", f"{max_diff_jit_aot:.2e}"),
    ],
    title="AOT vs JIT latency",
    aligns=["left", "right"],
)

Note that lower() takes a jax.ShapeDtypeStruct, not real data. You never need the input array to compile — only its shape and dtype.

You should see an AOT vs JIT latency table in which the AOT first execution is close to the AOT steady-state number rather than to the JIT first call and a Max |JIT − AOT| value that is effectively zero.

Inspect the StableHLO IR

lowered.as_text() shows the StableHLO program that will run on the device. This is the same intermediate representation XLA uses across GPUs, TPUs, and CPUs, and it is useful for debugging, performance analysis, and understanding what the compiler actually sees.

StableHLO text can be very large, and a single line may contain a long constant or attribute. To avoid Jupyter's IOPub data-rate limit, the next code saves the full IR to disk and prints only a bounded preview.

hlo_text = lowered.as_text()

hlo_path = pathlib.Path("/tmp/jax-course/l8-stablehlo.mlir")
hlo_path.parent.mkdir(parents=True, exist_ok=True)
hlo_path.write_text(hlo_text)

MAX_LINES = 20
MAX_CHARS_PER_LINE = 160

lines = hlo_text.splitlines()
preview_lines = []
for line in lines[:MAX_LINES]:
    if len(line) > MAX_CHARS_PER_LINE:
        preview_lines.append(line[:MAX_CHARS_PER_LINE] + " ... [line truncated]")
    else:
        preview_lines.append(line)

print(f"StableHLO program: {len(lines):,} lines, {len(hlo_text):,} chars")
print(f"Full StableHLO saved to: {hlo_path}")
print("=" * 60)
print("\n".join(preview_lines))
print(
    f"\n... ({max(len(lines) - MAX_LINES, 0):,} more lines; "
    "long lines are truncated in this preview)"
)

You should see the program's line and character counts, the path of the saved .mlir file, and the first 20 lines of the IR, followed by a count of how many lines were left out.

6. Measure batched inference throughput

GPUs are most efficient when processing many inputs at once. This step measures how forward-pass throughput scales with batch size. The model function does not change, but each new input shape needs its own lowered and compiled executable, so the loop compiles once per batch size.

batch_sizes = [1, 4, 16, 64]
results = []

for bs in batch_sizes:
    tokens_batch = jnp.zeros((bs, MAX_SEQ_LEN), dtype=jnp.int32)

    # Compile for this batch size
    lowered_bs = predict_closed.lower(
        jax.ShapeDtypeStruct((bs, MAX_SEQ_LEN), jnp.int32),
    )
    compiled_bs = lowered_bs.compile()

    # Warmup
    block_tree(compiled_bs(tokens_batch))

    # Measure
    times_bs = []
    for _ in range(50):
        start = time.perf_counter()
        block_tree(compiled_bs(tokens_batch))
        times_bs.append((time.perf_counter() - start) * 1000)

    avg_bs = np.mean(times_bs)
    tokens_per_sec = (bs * MAX_SEQ_LEN) / (avg_bs / 1000)
    results.append((bs, f"{avg_bs:.2f}", f"{tokens_per_sec:,.0f}"))

show_table(
    ["Batch size", "Latency (ms)", "Tokens/sec"],
    results,
    title="Batched inference throughput",
    aligns=["right", "right", "right"],
)

You should see one row per batch size, each with an average latency and a tokens/sec figure. Read the two columns together: latency per call and throughput per second answer different questions, and a serving system usually has to trade one against the other.

7. Path 3: export a portable artifact with jax.export

jax.export exports a jitted JAX function into an Exported object containing StableHLO plus the metadata needed to call it from another JAX process. The serialized bytes can be:

Saved to disk and loaded in another process
Called without the original Python model source code
Exported for the current platform by default, or for explicit platforms with the platforms=[...] argument

This is the JAX-native deployment path — no TensorFlow dependency required.

from jax import export

# Export the closure-based function
exported = export.export(predict_closed)(
    jax.ShapeDtypeStruct((1, MAX_SEQ_LEN), jnp.int32),
)

print(f"Exported function: {exported.fun_name}")
print(f"Input shapes:  {exported.in_avals}")
print(f"Output shapes: {exported.out_avals}")
print(f"Exported platforms: {exported.platforms}")

# Serialize to bytes
blob = exported.serialize()
export_path = pathlib.Path("/tmp/jax-course/exports")
export_path.mkdir(parents=True, exist_ok=True)

export_file = export_path / "tiny_transformer_jax_export.bin"
export_file.write_bytes(blob)
print()
print(f"Serialized to {export_file} ({len(blob):,} bytes, {len(blob) / 1024:.0f} KB)")

# Deserialize and call
rehydrated = export.deserialize(export_file.read_bytes())

test_input = jnp.zeros((1, MAX_SEQ_LEN), dtype=jnp.int32)
logits_exported = block_tree(rehydrated.call(test_input))
print()
print(f"✅ Deserialized call succeeded — output shape: {logits_exported.shape}")

# Verify outputs match
diff_export = float(jnp.max(jnp.abs(logits_aot - logits_exported)))
print(f"Max difference from AOT: {diff_export:.2e}")

The rehydrated object never saw TinyTransformer. It was rebuilt from bytes on disk and still returns the same logits, which is the whole point of a portable artifact.

You should see the exported function name, the input and output avals, the exported platform list, a byte count for the serialized blob, an output shape from the deserialized call, and a max difference from the AOT logits that is effectively zero.

8. Path 4: convert to a TensorFlow SavedModel with jax2tf

If your serving infrastructure uses TensorFlow such as TF Serving or TFX pipelines, you can convert the JAX function to a TF SavedModel. jax2tf still lives under jax.experimental, but it is the standard JAX-to-TensorFlow interop path.

Native serialization is the default in current JAX releases, so the conversion embeds lowered StableHLO in the TensorFlow graph without passing native_serialization=True.

One platform detail matters here. This codelab runs JAX on CUDA, but the TensorFlow runtime in the JAX container may execute the SavedModel on CPU. To avoid a CUDA-exported module being called by TensorFlow on CPU, the code below exports the jax2tf module for ("cpu",). If you serve with TensorFlow on GPU, export for ("cuda",) and use a TensorFlow runtime with GPU/XLA support.

from jax.experimental import jax2tf
import tensorflow as tf
import shutil

# Capture model_state as a closure
def predict_for_tf(tokens):
    m = nnx.merge(graphdef, model_state)
    return m(tokens)

TF_EXPORT_PLATFORMS = ("cpu",)
tf_predict = jax2tf.convert(
    predict_for_tf,
    native_serialization_platforms=TF_EXPORT_PLATFORMS,
)

# Wrap in a tf.Module for SavedModel export
module = tf.Module()
module.predict = tf.function(
    tf_predict,
    input_signature=[tf.TensorSpec(shape=(1, MAX_SEQ_LEN), dtype=tf.int32)],
    autograph=False,
)

# TF Serving expects a versioned model directory
savedmodel_base_dir = pathlib.Path("/tmp/jax-course/exports") / "tiny_transformer_savedmodel"
savedmodel_dir = savedmodel_base_dir / "1"
if savedmodel_base_dir.exists():
    shutil.rmtree(savedmodel_base_dir)

tf.saved_model.save(module, str(savedmodel_dir))
print(f"✅ SavedModel saved to {savedmodel_dir}")
print(f"Exported for TensorFlow platform(s): {TF_EXPORT_PLATFORMS}")

# Verify the SavedModel path produces the same logits as the JAX path
with tf.device("/CPU:0"):
    tf_logits = module.predict(tf.zeros((1, MAX_SEQ_LEN), dtype=tf.int32))
diff_tf = np.max(np.abs(np.asarray(tf_logits) - np.asarray(logits_aot)))
print(f"Max difference from AOT: {diff_tf:.2e}")

# List saved files
for dirpath, _, filenames in os.walk(savedmodel_base_dir):
    for f in filenames:
        full = os.path.join(dirpath, f)
        size = os.path.getsize(full)
        print(f"  {os.path.relpath(full, savedmodel_base_dir):44s} {size:>10,} bytes")
print()
print(
    f"To serve: docker run -p 8501:8501 "
    f"--mount type=bind,source={savedmodel_base_dir},target=/models/transformer "
    "-e MODEL_NAME=transformer tensorflow/serving "
    "--xla_cpu_compilation_enabled=true"
)

You should see the SavedModel directory path, the export platform tuple, a max difference from the AOT logits that is effectively zero, and a listing of the files written under the versioned directory.

The last thing the code prints is a docker run ... tensorflow/serving command. It is illustrative and it shows how you would point TensorFlow Serving at the versioned directory you just created. You do not run it in this codelab.

9. Compare the four serving paths

Each path started from the same trained checkpoint and produced the same predictions. What differs is portability, dependencies, deployment target, and whether the artifact can be used outside the original Python process.

show_table(
    ["Path", "Format", "Dependencies", "Serving target", "Portable"],
    [
        ("jax.jit", "Cached in-process executable", "JAX", "Python server", "No"),
        ("AOT compile", "Compiled in-process executable", "JAX", "Python server startup / warmup", "No"),
        ("jax.export", "Serialized JAX export", "JAX runtime", "Compatible runtime for exported platform(s)", "Yes"),
        ("jax2tf", "TF SavedModel", "TensorFlow", "TF Serving", "Yes"),
    ],
    title="Serving path comparison",
)

# Show file sizes
sizes = []
export_size = os.path.getsize(export_file)
sizes.append(("jax.export", f"{export_size:,} bytes", f"{export_size / 1024:.0f} KB"))

sm_size = sum(
    os.path.getsize(os.path.join(dirpath, f))
    for dirpath, _, filenames in os.walk(savedmodel_base_dir)
    for f in filenames
)
sizes.append(("jax2tf SavedModel", f"{sm_size:,} bytes", f"{sm_size / 1024:.0f} KB"))

show_table(
    ["Export", "Size (bytes)", "Size (KB)"],
    sizes,
    title="Export file sizes",
    aligns=["left", "right", "right"],
)

You should see the Serving path comparison table followed by an Export file sizes table listing the on-disk size of the jax.export blob and of the SavedModel directory.

Choose based on your deployment constraints, not on model architecture. If JAX is already running in your serving process, jax.jit with a startup warmup or AOT compilation is the shortest path. If the artifact has to leave that process, jax.export keeps you inside JAX and jax2tf hands you over to the TensorFlow ecosystem.

10. Clean up

Everything this codelab wrote lives under /tmp on the Pod, which disappears with the Pod. Copy off anything you want to keep first, running this from Cloud Shell:

kubectl cp jax-jupyter:/tmp/jax-course/exports ./jax-course-exports

Delete the Jupyter workload, including the LoadBalancer and the persistent volume:

kubectl delete -f deploy/jupyter.yaml

Destroy the cluster, node pool, VPC, and service account:

cd terraform
terraform destroy

Type yes when prompted, then confirm nothing is left behind:

gcloud container clusters list
gcloud compute instances list

Both should be empty for this project. If you created a project just for this series, you can instead delete the whole project from the Cloud console.

11. Congratulations

You took a trained transformer checkpoint out of the notebook and turned it into four serving-ready artifacts on an NVIDIA L4 GPU.

What you've learned

jax.jit is the simplest path — wrap and call. The first call compiles; subsequent calls are fast. Good for Python-based serving (FastAPI, Flask) where JAX is already installed.
AOT compilation (lower() then compile()) separates compilation from execution. Compile once at startup, then serve without first-request compilation latency. lowered.as_text() shows the StableHLO IR for debugging and performance analysis.
jax.export serializes a jitted function to a JAX-native artifact containing StableHLO and calling metadata. The resulting file can be loaded and called by a compatible JAX runtime without the original model code. By default it exports for the current platform; use platforms=[...] when you need an explicit target.
jax2tf converts the function to a TensorFlow SavedModel. Use it when your serving infrastructure is TensorFlow-based. Native serialization is the default in current JAX releases; set native_serialization_platforms to match where TensorFlow will execute the model.
How to restore an Orbax checkpoint into a rebuilt architecture using jax.ShapeDtypeStruct targets, and how to move parameters onto a single device with jax.device_put
Why forward-pass tokens/sec is not autoregressive generation throughput, and why each new input shape needs its own compiled executable

Course recap

Over eight labs, you have:

L1–L3: Set up JAX, learned jit compilation, and profiled GPU execution
L4: Built a training loop from scratch — optimizer, loss, gradient updates
L5: Explored attention mechanisms — naive, SDPA, cuDNN fused attention
L6: Scaled to multiple GPUs with data parallelism — Mesh, NamedSharding, data-parallel sharding
L7: Combined everything into a transformer language model — Flax NNX, Orbax, generation
L8: Prepared the trained model for production — JIT, AOT, jax.export, jax2tf

Next steps

You now have the core workflow for building with JAX on GPUs. The best next step is to turn these pieces into larger, messier projects:

Add KV-cache decoding and batched generation, so you measure real generation throughput rather than forward-pass throughput
Train on a real tokenizer and dataset instead of the toy vocabulary used here
Experiment with mixed precision and quantization, and measure rather than assume the speedup
Scale beyond data parallelism into model and tensor parallelism
Build a small serving stack around one of today's artifacts, with monitoring and load tests

From here the focus shifts from learning individual JAX features to making engineering choices: how to keep models fast, reproducible, memory-efficient, debuggable, and deployable.