Train a model on GPU with JAX, Optax, and Fashion-MNIST

1. Introduction

In the previous codelabs, you verified that JAX can see the GPU, learned how jax.jit traces and compiles a function, and used the profiler to see what the GPU was actually doing. Now the pieces come together into the most common JAX workflow: a training loop.

In this one you train a small MLP on Fashion-MNIST, a real image-classification dataset with 60,000 training examples and 10,000 test examples. Each example is a 28x28 grayscale image of a clothing item. By the end you will have a compiled Optax training step, throughput numbers you can trust, and a confusion matrix that connects those numbers back to real images.

What you'll do

• Download Fashion-MNIST and reshape it into fixed-size batches that live on the GPU
• Build a small MLP as a PyTree of JAX arrays and write a scalar loss function
• Compute gradients with jax.grad and jax.value_and_grad, then compile the step with jax.jit
• Replace hand-written SGD with an Optax AdamW optimizer and run a short training loop
• Measure throughput in examples/sec and relate it to tokens/sec
• Evaluate the model, plot a confusion matrix, and compare float32 with bfloat16 compute

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 3, or an equivalent CUDA-enabled JAX GPU environment
• Internet access from the Pod for the first Fashion-MNIST download

Estimated time to complete: 60 minutes.

The training-step mental model

A JAX training step is a pure function: it takes arrays in, returns new arrays out, and does not mutate the old parameters in place. Every piece below has a beginner check you can apply when something goes wrong.

Piece	What it does	Beginner check
params	Model weights stored as a PyTree	Same tree structure as grads
batch	Images and labels	Same shapes every step to avoid recompilation
loss_fn	Forward pass plus scalar loss	jax.grad needs a scalar loss
jax.value_and_grad	Computes loss and gradients together	Gradients match parameter shapes
optimizer.update	Converts gradients into updates	Adam stores optimizer state
optax.apply_updates	Produces the next parameters	Parameters are immutable, so return the new tree

Those pieces run in a fixed four-stage cycle, from forward step to gradients update and repeat.

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 10 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 optax matplotlib

Both packages usually ship in the NVIDIA JAX container, so this pip command is normally a no-op.

Set up and verify the GPU

Import JAX, Optax, and a few helpers. This cell also verifies that the default backend is a GPU.

import gzip
import hashlib
import html
import math
import pathlib
import struct
import time
import urllib.request
from functools import partial

from IPython.display import HTML, display
import matplotlib.pyplot as plt
import numpy as np

import jax
import jax.numpy as jnp

try:
    import optax
except ModuleNotFoundError as exc:
    raise ModuleNotFoundError(
        "This lesson requires Optax. Install it with: pip install optax"
    ) from exc


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

print(f"JAX version:     {jax.__version__}")
print(f"Optax version:   {getattr(optax, '__version__', 'unknown')}")
print(f"Default backend: {jax.default_backend()}")
print(f"Devices:         {devices}")

assert gpu_devices, f"This lesson assumes a GPU backend. Available devices: {devices}"
print(f"Using GPU:       {device}")


CLASS_NAMES = np.array([
    "T-shirt/top",
    "Trouser",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
])


def block_tree(tree):
    """Wait until a PyTree of JAX arrays is ready on device."""
    return jax.block_until_ready(tree)


def tree_l2_norm(tree):
    """L2 norm of all leaves in a PyTree treated as one long vector."""
    leaves = jax.tree_util.tree_leaves(tree)
    return jnp.sqrt(sum(jnp.sum(jnp.square(x)) for x in leaves))


def count_params(params):
    """Total number of scalar values across all leaves of a parameter PyTree."""
    return sum(x.size for x in jax.tree_util.tree_leaves(params))


def show_table(headers, rows, title=None, aligns=None):
    """Render rows as an HTML table; `aligns` is an optional per-column list of "left"/"right"/"center"."""
    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("".join(parts)))


def show_bars(rows, title, unit="", lower_is_better=False):
    """Render (label, value) pairs as a horizontal bar chart in HTML. Scales bars to the largest value."""
    max_value = max(float(value) for _, value in rows) or 1.0
    color = "#1a7f37" if not lower_is_better else "#0969da"
    parts = ["<div style='font-family: Arial, sans-serif; max-width: 760px;'>"]
    parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    for label, value in rows:
        width = max(3, 100 * float(value) / max_value)
        parts.append(
            "<div style='display:grid; grid-template-columns: 190px 1fr 130px; gap: 8px; "
            "align-items:center; margin: 6px 0;'>"
            f"<div style='font-size:13px;'>{html.escape(str(label))}</div>"
            "<div style='background:#f6f8fa; border-radius:6px; overflow:hidden; height:22px;'>"
            f"<div style='height:22px; width:{width:.1f}%; background:{color};'></div></div>"
            f"<div style='font-size:13px; font-variant-numeric: tabular-nums;'>{float(value):,.1f} {html.escape(unit)}</div>"
            "</div>"
        )
    parts.append(
        f"<div style='font-size:12px; color:#57606a;'>"
        f"{'Lower' if lower_is_better else 'Higher'} is better.</div></div>"
    )
    display(HTML("".join(parts)))

You should see a JAX version, an Optax version, gpu as the default backend, and a list of CUDA devices, followed by the GPU the rest of the codelab will use. The show_table and show_bars helpers render the result tables and bar charts you will see in later steps.

3. Load and inspect Fashion-MNIST

Before you can train anything you need data on the GPU in a shape that will not change from step to step. This step downloads Fashion-MNIST, looks at it, and turns it into fixed-size batches that live on the device.

Download the dataset

Fashion-MNIST uses the same IDX file format as MNIST. The helper functions below download the compressed files, verify their checksums, and parse images and labels into NumPy arrays.

This repository is the primary source. The dataset files are cached locally, so this cell should be quick after the first run.

DATA_DIR = pathlib.Path.home() / ".cache" / "jax-course" / "fashion-mnist"
DATA_DIR.mkdir(parents=True, exist_ok=True)

FILES = {
    "train-images-idx3-ubyte.gz": "8d4fb7e6c68d591d4c3dfef9ec88bf0d",
    "train-labels-idx1-ubyte.gz": "25c81989df183df01b3e8a0aad5dffbe",
    "t10k-images-idx3-ubyte.gz": "bef4ecab320f06d8554ea6380940ec79",
    "t10k-labels-idx1-ubyte.gz": "bb300cfdad3c16e7a12a480ee83cd310",
}

PRIMARY_BASE_URL = "https://github.com/zalandoresearch/fashion-mnist/raw/master/data/fashion"

def md5sum(path):
    """Stream `path` in 1 MB chunks and return its MD5 hex digest."""
    digest = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def download_if_needed(filename, expected_md5):
    """Download `filename` if missing or its MD5 doesn't match."""
    path = DATA_DIR / filename
    if path.exists() and md5sum(path) == expected_md5:
        print(f"Using cached {filename}")
        return path

    urls = [f"{PRIMARY_BASE_URL}/{filename}"]
    last_error = None
    for url in urls:
        try:
            print(f"Downloading {filename} from {url}")
            urllib.request.urlretrieve(url, path)
            actual_md5 = md5sum(path)
            if actual_md5 != expected_md5:
                raise ValueError(f"MD5 mismatch: expected {expected_md5}, got {actual_md5}")
            return path
        except Exception as exc:
            last_error = exc
            if path.exists():
                path.unlink()
            print(f"  failed: {exc}")

    raise RuntimeError(f"Could not download {filename}") from last_error


def read_idx_images(path):
    """Parse the Fashion-MNIST IDX-3 image file at `path` and return a (N, rows, cols) uint8 array."""
    with gzip.open(path, "rb") as f:
        magic, num_images, rows, cols = struct.unpack(">IIII", f.read(16))
        assert magic == 2051, f"Unexpected image magic number {magic} in {path}"
        data = np.frombuffer(f.read(), dtype=np.uint8)
    return data.reshape(num_images, rows, cols)


def read_idx_labels(path):
    """Parse the IDX-1 label file at `path` and return a 1-D uint8 array of class indices."""
    with gzip.open(path, "rb") as f:
        magic, num_labels = struct.unpack(">II", f.read(8))
        assert magic == 2049, f"Unexpected label magic number {magic} in {path}"
        data = np.frombuffer(f.read(), dtype=np.uint8)
    return data.reshape(num_labels)


paths = {name: download_if_needed(name, checksum) for name, checksum in FILES.items()}

train_images = read_idx_images(paths["train-images-idx3-ubyte.gz"])
train_labels = read_idx_labels(paths["train-labels-idx1-ubyte.gz"])
test_images = read_idx_images(paths["t10k-images-idx3-ubyte.gz"])
test_labels = read_idx_labels(paths["t10k-labels-idx1-ubyte.gz"])

show_table(
    ["Split", "Images", "Image shape", "Labels"],
    [
        ("train", f"{len(train_images):,}", train_images.shape[1:], f"{len(train_labels):,}"),
        ("test", f"{len(test_images):,}", test_images.shape[1:], f"{len(test_labels):,}"),
    ],
    title="Fashion-MNIST loaded from IDX files",
)

The table should report 60,000 training images and 10,000 test images, each with shape (28, 28), and the same number of labels as images in both splits.

Look at a few examples

Before training, always look at a few examples. This catches many boring but expensive bugs: wrong labels, wrong image orientation, wrong scaling, or accidentally loading the wrong dataset.

fig, axes = plt.subplots(2, 5, figsize=(10, 4))
for label, ax in enumerate(axes.flat):
    idx = np.flatnonzero(train_labels == label)[0]
    ax.imshow(train_images[idx], cmap="gray")
    ax.set_title(CLASS_NAMES[label], fontsize=10)
    ax.axis("off")
fig.suptitle("One Fashion-MNIST example per class")
fig.tight_layout()
plt.show()

You should see a two-by-five grid with one recognizable clothing item per class, and each title matching the image below it.

Prepare fixed-size GPU batches

Now that you checked data, you can get them into the shape the training loop wants. JAX is more efficient when every training step receives the same shapes and dtypes.

The cell normalizes pixels to [0, 1], flattens each 28x28 image into a 784-value vector, shuffles the training set once, and reshapes the data into fixed-size batches. The batches are moved to the GPU once with jax.device_put, so the training loop only indexes arrays that are already on the device.

TRAIN_EXAMPLES = 60_000
TEST_EXAMPLES = 10_000
BATCH_SIZE = 512
INPUT_DIM = 28 * 28
NUM_CLASSES = 10


def prepare_images(images):
    """Cast uint8 images to float32 in [0, 1] and flatten each one into a 1-D feature vector."""
    images = images.astype(np.float32) / 255.0
    return images.reshape(images.shape[0], -1)


rng = np.random.default_rng(0)
train_perm = rng.permutation(len(train_images))[:TRAIN_EXAMPLES]

x_train = prepare_images(train_images[train_perm])
y_train = train_labels[train_perm].astype(np.int32)
x_test = prepare_images(test_images[:TEST_EXAMPLES])
y_test = test_labels[:TEST_EXAMPLES].astype(np.int32)


def make_fixed_batches(x, y, batch_size):
    """Trim trailing examples that don't fill a batch, reshape, and move to device."""
    usable = (len(x) // batch_size) * batch_size
    x = x[:usable].reshape(usable // batch_size, batch_size, x.shape[-1])
    y = y[:usable].reshape(usable // batch_size, batch_size)
    return jax.device_put(jnp.asarray(x), device), jax.device_put(jnp.asarray(y), device)


x_train_batches, y_train_batches = make_fixed_batches(x_train, y_train, BATCH_SIZE)
x_test_batches, y_test_batches = make_fixed_batches(x_test, y_test, BATCH_SIZE)
first_batch = (x_train_batches[0], y_train_batches[0])

show_table(
    ["Array", "Shape", "Dtype", "Devices"],
    [
        ("x_train_batches", x_train_batches.shape, x_train_batches.dtype, x_train_batches.devices()),
        ("y_train_batches", y_train_batches.shape, y_train_batches.dtype, y_train_batches.devices()),
        ("x_test_batches", x_test_batches.shape, x_test_batches.dtype, x_test_batches.devices()),
        ("y_test_batches", y_test_batches.shape, y_test_batches.dtype, y_test_batches.devices()),
    ],
    title="Fixed-size batches on GPU",
)

In the table, every array should be float32 or int32, every shape should end in 784 for images and 512 for labels, and the Devices column should show the CUDA device you selected in the setup cell.

4. Define the model and take one gradient step

With batches on the device, you need two things: a model that turns a batch into logits, and a scalar loss you can differentiate. This step builds both, then takes a single gradient step by hand so you can see exactly what jax.grad returns.

Define the model and loss

The model is a small MLP. A convolutional model would usually be better for images, but an MLP keeps the training-step mechanics visible: matrix multiply, nonlinearity, matrix multiply, loss, gradients, optimizer update.

The forward pass casts weights and activations to compute_dtype, then casts logits back to float32 before the loss. For now the compute dtype is float32 but later you will switch it to bfloat16 and measure what changes.

HIDDEN1 = 256
HIDDEN2 = 128
LEARNING_RATE = 3e-3


def init_mlp_params(key, input_dim=INPUT_DIM, hidden1=HIDDEN1, hidden2=HIDDEN2, num_classes=NUM_CLASSES):
    """Initialize a 3-layer MLP with He-style weight scaling and zero biases."""
    k1, k2, k3 = jax.random.split(key, 3)
    return {
        "w1": jax.random.normal(k1, (input_dim, hidden1), dtype=jnp.float32) * math.sqrt(2.0 / input_dim),
        "b1": jnp.zeros((hidden1,), dtype=jnp.float32),
        "w2": jax.random.normal(k2, (hidden1, hidden2), dtype=jnp.float32) * math.sqrt(2.0 / hidden1),
        "b2": jnp.zeros((hidden2,), dtype=jnp.float32),
        "w3": jax.random.normal(k3, (hidden2, num_classes), dtype=jnp.float32) * math.sqrt(2.0 / hidden2),
        "b3": jnp.zeros((num_classes,), dtype=jnp.float32),
    }


def mlp(params, x, compute_dtype=jnp.float32):
    """Forward pass: cast inputs/params to `compute_dtype`, two GELU hidden layers, then cast logits back to float32."""
    x = x.astype(compute_dtype)
    w1 = params["w1"].astype(compute_dtype)
    b1 = params["b1"].astype(compute_dtype)
    w2 = params["w2"].astype(compute_dtype)
    b2 = params["b2"].astype(compute_dtype)
    w3 = params["w3"].astype(compute_dtype)
    b3 = params["b3"].astype(compute_dtype)

    x = jax.nn.gelu(x @ w1 + b1)
    x = jax.nn.gelu(x @ w2 + b2)
    logits = x @ w3 + b3
    return logits.astype(jnp.float32)


def cross_entropy_loss(params, batch, compute_dtype=jnp.float32):
    """Scalar softmax cross-entropy loss."""
    x, y = batch
    logits = mlp(params, x, compute_dtype=compute_dtype)
    return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()


def loss_with_metrics(params, batch, compute_dtype=jnp.float32):
    """Same loss, but also returns batch accuracy in an aux dict."""
    x, y = batch
    logits = mlp(params, x, compute_dtype=compute_dtype)
    loss = optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()
    accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == y)
    return loss, {"accuracy": accuracy}


params = init_mlp_params(jax.random.key(1))
params = jax.device_put(params, device)

rows = []
for name, value in params.items():
    rows.append((name, value.shape, value.dtype, value.devices()))
show_table(["Parameter", "Shape", "Dtype", "Devices"], rows, title=f"MLP parameters: {count_params(params):,} trainable values")

The table lists six leaves — three weight matrices and three bias vectors — all float32 and all on the GPU. The title reports the total number of trainable values.

Compute one gradient step by hand

jax.grad and jax.value_and_grad differentiate with respect to the first argument by default. Here the first argument is params, so the gradient has the same tree structure as the parameter dictionary.

Use jax.grad when you only need gradients. Use jax.value_and_grad in a training step when you want the loss and gradients from the same forward and backward pass.

The code below is the simplest possible training step: compute the loss, compute gradients, subtract a scaled gradient from every parameter, and return a new parameter tree.

def sgd_step(params, batch):
    """One un-jitted SGD update and returns (new_params, loss, grads)."""
    loss, grads = jax.value_and_grad(cross_entropy_loss)(params, batch)
    new_params = jax.tree.map(lambda p, g: p - LEARNING_RATE * g, params, grads)
    return new_params, loss, grads


grads_only = jax.grad(cross_entropy_loss)(params, first_batch)
loss_value, grads = jax.value_and_grad(cross_entropy_loss)(params, first_batch)
loss_value, grads, grads_only = block_tree((loss_value, grads, grads_only))

rows = []
for name in params:
    rows.append((name, params[name].shape, grads[name].shape, grads[name].dtype))
show_table(["Leaf", "Param shape", "Grad shape", "Grad dtype"], rows, title="Gradient tree matches the parameter tree")

grad_difference = tree_l2_norm(jax.tree.map(lambda a, b: a - b, grads, grads_only))

print(f"loss before update: {float(loss_value):.4f}")
print(f"gradient L2 norm:   {float(tree_l2_norm(grads)):.4f}")
print(f"grad vs value_and_grad difference: {float(grad_difference):.6f}")

params_after_one, loss_after_one, _ = sgd_step(params, first_batch)
params_after_one, loss_after_one = block_tree((params_after_one, loss_after_one))
print(f"loss used for one SGD update: {float(loss_after_one):.4f}")

Three things to check in the output. Every row has a Grad shape identical to its Param shape, the grad vs value_and_grad difference is zero or very close to it, because both transformations compute the same derivative and the last printed loss is the loss computed before the update, on the same batch.

sgd_step returns a new parameter tree rather than modifying params, which is what "pure function" means in practice.

5. Compile the training step with jax.jit

The hand-written SGD step is correct, but it is not the way you want to run a GPU training loop. Without jit, Python keeps dispatching many small operations. With jit, JAX traces the whole step once and XLA compiles it into an executable for this batch shape and dtype.

The first jitted call includes compilation. The timing below warms up once, then measures cached execution.

@jax.jit
def sgd_step_jit(params, batch):
    loss, grads = jax.value_and_grad(cross_entropy_loss)(params, batch)
    new_params = jax.tree.map(lambda p, g: p - LEARNING_RATE * g, params, grads)
    return new_params, loss


def batch_at(step):
    """Pick batch `step % num_batches`."""
    i = step % x_train_batches.shape[0]
    return x_train_batches[i], y_train_batches[i]


def time_loop(step_fn, params, steps):
    """Run `step_fn` for `steps` iterations and time it."""
    start = time.perf_counter()
    loss = None
    for step in range(steps):
        params, loss = step_fn(params, batch_at(step))
    params, loss = block_tree((params, loss))
    elapsed = time.perf_counter() - start
    return params, loss, elapsed


params_warm, loss_warm = sgd_step_jit(params, first_batch)
block_tree((params_warm, loss_warm))

EAGER_STEPS = 20
JIT_STEPS = 100

# sgd_step returns (params, loss, grads).
_, eager_loss, eager_elapsed = time_loop(lambda p, b: sgd_step(p, b)[:2], params, EAGER_STEPS)
_, jit_loss, jit_elapsed = time_loop(sgd_step_jit, params, JIT_STEPS)

eager_rate = EAGER_STEPS * BATCH_SIZE / eager_elapsed
jit_rate = JIT_STEPS * BATCH_SIZE / jit_elapsed

show_table(
    ["Mode", "Steps", "Final loss", "Elapsed seconds", "Examples/sec"],
    [
        ("Python dispatch", EAGER_STEPS, f"{float(eager_loss):.4f}", f"{eager_elapsed:.3f}", f"{eager_rate:,.0f}"),
        ("jitted step", JIT_STEPS, f"{float(jit_loss):.4f}", f"{jit_elapsed:.3f}", f"{jit_rate:,.0f}"),
    ],
    title="Cached training-step throughput",
    aligns=["left", "right", "right", "right", "right"],
)
show_bars([("Python dispatch", eager_rate), ("jitted step", jit_rate)], "Examples per second", "examples/s")

You should see two rows and a two-bar chart, with the jitted step reaching a higher examples/sec than Python dispatch. The exact ratio depends on your GPU and on the batch shape, but the ordering is the point: the compiled step does the same work with far less per-operation overhead.

6. Train with an Optax optimizer

Raw SGD is enough to demonstrate gradients, but real training loops want momentum, weight decay, and schedules. This step swaps in Optax, runs a real training loop, and turns the result into a throughput number.

Set up the Optax optimizer

Optax, a gradient processing and optimization library for JAX, gives you composable optimizers such as SGD with momentum, Adam, AdamW, gradient clipping, and learning-rate schedules. An Optax optimizer has two important methods: optimizer.init(params) creates the optimizer state, such as Adam's momentum buffers, and optimizer.update(grads, opt_state, params) converts gradients into updates and returns the next optimizer state. Then optax.apply_updates(params, updates) returns the next parameter tree.

optimizer = optax.adamw(learning_rate=LEARNING_RATE, weight_decay=1e-4)
opt_state = optimizer.init(params)


@jax.jit
def train_step(params, opt_state, batch):
    (loss, metrics), grads = jax.value_and_grad(loss_with_metrics, has_aux=True)(params, batch)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    metrics = {
        "loss": loss,
        "accuracy": metrics["accuracy"],
        "grad_norm": optax.global_norm(grads),
    }
    return params, opt_state, metrics


params_opt = init_mlp_params(jax.random.key(2))
params_opt = jax.device_put(params_opt, device)
opt_state = optimizer.init(params_opt)

params_opt, opt_state, metrics = train_step(params_opt, opt_state, first_batch)
params_opt, opt_state, metrics = block_tree((params_opt, opt_state, metrics))

show_table(
    ["Metric", "Value"],
    [
        ("loss", f"{float(metrics['loss']):.4f}"),
        ("accuracy", f"{100 * float(metrics['accuracy']):.1f}%"),
        ("gradient L2 norm", f"{float(metrics['grad_norm']):.4f}"),
    ],
    title="One compiled Optax training step",
    aligns=["left", "right"],
)

The table shows the loss, accuracy, and gradient norm from a single compiled step on a freshly initialized model, so the accuracy should be near chance level for ten classes.

Run a short training loop

One step works, so you can repeat it. The loop below calls the same compiled train_step on same-shaped batches. That is the core pattern for JAX training.

Notice the logging pattern. Metrics stay on the device during the step, and are converted to Python floats only every few steps. Pulling a loss to Python every iteration is convenient, but it also synchronizes the host with the GPU every iteration.

def train_many_steps(params, opt_state, steps=400, log_every=25):
    """Run a sized training loop, log a metric snapshot every `log_every` steps, and return history plus final state."""
    history = []
    start = time.perf_counter()
    metrics = None

    for step in range(steps):
        params, opt_state, metrics = train_step(params, opt_state, batch_at(step))

        if step % log_every == 0 or step == steps - 1:
            metrics = block_tree(metrics)
            history.append(
                {
                    "step": step,
                    "loss": float(metrics["loss"]),
                    "accuracy": float(metrics["accuracy"]),
                    "grad_norm": float(metrics["grad_norm"]),
                }
            )

    params, opt_state, metrics = block_tree((params, opt_state, metrics))
    elapsed = time.perf_counter() - start
    return params, opt_state, history, elapsed, metrics


params_train = init_mlp_params(jax.random.key(3))
params_train = jax.device_put(params_train, device)
# Fresh optimizer state, paired with this fresh `params_train`.
opt_state = optimizer.init(params_train)

params_train, opt_state, _ = train_step(params_train, opt_state, first_batch)
block_tree((params_train, opt_state))

TRAIN_STEPS = 400
params_train, opt_state, history, elapsed, final_metrics = train_many_steps(
    params_train, opt_state, steps=TRAIN_STEPS, log_every=25
)
examples_per_sec = TRAIN_STEPS * BATCH_SIZE / elapsed

show_table(
    ["Step", "Loss", "Accuracy", "Grad norm"],
    [(h["step"], f"{h['loss']:.4f}", f"{100*h['accuracy']:.1f}%", f"{h['grad_norm']:.3f}") for h in history],
    title=f"Training metrics, {examples_per_sec:,.0f} examples/sec",
    aligns=["right", "right", "right", "right"],
)

steps = [h["step"] for h in history]
losses = [h["loss"] for h in history]
accuracies = [h["accuracy"] for h in history]

fig, ax1 = plt.subplots(figsize=(8, 4))
ax1.plot(steps, losses, marker="o", color="#0969da", label="loss")
ax1.set_xlabel("step")
ax1.set_ylabel("loss", color="#0969da")
ax1.tick_params(axis="y", labelcolor="#0969da")
ax1.grid(True, alpha=0.25)

ax2 = ax1.twinx()
ax2.plot(steps, accuracies, marker="s", color="#1a7f37", label="accuracy")
ax2.set_ylabel("batch accuracy", color="#1a7f37")
ax2.tick_params(axis="y", labelcolor="#1a7f37")
ax2.set_ylim(0.0, 1.0)

fig.suptitle("Fashion-MNIST training curve")
fig.tight_layout()
plt.show()

You should see a metrics table with one row per logged step and a two-axis plot in which the loss curve falls and the batch-accuracy curve rises over the 400 steps.

7. Evaluate the trained model

Training-batch accuracy tells you whether the optimizer is making progress. This step checks whether the model learned something that generalizes, then makes the errors concrete by looking at individual predictions and a confusion matrix.

Evaluate on held-out test data

Test accuracy is a better initial check that the model learned something useful rather than only memorizing a few batches.

The evaluation function uses jax.vmap over fixed-size test batches. That keeps the code compact and lets JAX run the same prediction logic over all batches.

@jax.jit
def evaluate_batches(params, x_batches, y_batches):
    """Vmap `loss_with_metrics` over every (x, y) batch and return the mean loss and accuracy."""
    def eval_one_batch(x, y):
        loss, metrics = loss_with_metrics(params, (x, y))
        return loss, metrics["accuracy"]

    losses, accuracies = jax.vmap(eval_one_batch)(x_batches, y_batches)
    return {"loss": jnp.mean(losses), "accuracy": jnp.mean(accuracies)}


test_metrics = evaluate_batches(params_train, x_test_batches, y_test_batches)
test_metrics = block_tree(test_metrics)

show_table(
    ["Split", "Loss", "Accuracy", "Examples evaluated"],
    [
        ("train batch", f"{float(final_metrics['loss']):.4f}", f"{100 * float(final_metrics['accuracy']):.1f}%", BATCH_SIZE),
        ("test", f"{float(test_metrics['loss']):.4f}", f"{100 * float(test_metrics['accuracy']):.1f}%", int(np.prod(y_test_batches.shape))),
    ],
    title="Evaluation after the short training run",
    aligns=["left", "right", "right", "right"],
)

The test row should be in the same range as the final training batch, not dramatically worse. The Examples evaluated column shows the trimmed test set, not the full 10,000, because the batches you built earlier dropped the 272 leftovers.

Visualize predictions

Now it is time to look at predictions. Green titles are correct predictions. Red titles are mistakes.

The model is intentionally small and trained briefly, so a few mistakes are expected. And the mistakes are often between visually similar classes such as shirt, T-shirt/top, pullover, and coat.

@jax.jit
def predict(params, x, compute_dtype=jnp.float32):
    """Return the predicted class index (argmax of the logits) for each row of `x`."""
    logits = mlp(params, x, compute_dtype=compute_dtype)
    return jnp.argmax(logits, axis=-1)


rng = np.random.default_rng(7)
sample_count = 25
sample_indices = rng.choice(len(test_images), size=sample_count, replace=False)
sample_pixels = test_images[sample_indices]
sample_x = prepare_images(sample_pixels)
sample_y = test_labels[sample_indices].astype(np.int32)

sample_x_device = jax.device_put(jnp.asarray(sample_x), device)
sample_pred = np.asarray(block_tree(predict(params_train, sample_x_device)))

fig, axes = plt.subplots(5, 5, figsize=(10, 10))
for ax, image, true_label, pred_label in zip(axes.flat, sample_pixels, sample_y, sample_pred):
    correct = int(true_label) == int(pred_label)
    ax.imshow(image, cmap="gray")
    ax.set_title(
        f"pred: {CLASS_NAMES[pred_label]}\ntrue: {CLASS_NAMES[true_label]}",
        fontsize=9,
        color="#1a7f37" if correct else "#d1242f",
    )
    ax.axis("off")
fig.suptitle("Sample Fashion-MNIST predictions")
fig.tight_layout()
plt.show()

You should see a five-by-five grid dominated by green titles, with a handful of red ones. Look at the red titles: most of them should be confusions between clothing types that look alike in a 28x28 grayscale image.

Confusion matrix

The images we use are a small sample. A confusion matrix shows which classes the model mixes up across the whole test set. Rows are true labels and columns are predicted labels with the diagonal means the model is usually right.

@jax.jit
def predict_batches(params, x_batches):
    """Run `predict` over every batch with vmap; output shape is (num_batches, batch_size)."""
    return jax.vmap(lambda x: predict(params, x))(x_batches)


test_pred = np.asarray(block_tree(predict_batches(params_train, x_test_batches))).reshape(-1)
test_true = np.asarray(y_test_batches).reshape(-1)

confusion = np.zeros((NUM_CLASSES, NUM_CLASSES), dtype=np.int32)
np.add.at(confusion, (test_true, test_pred), 1)
confusion_percent = confusion / confusion.sum(axis=1, keepdims=True)

fig, ax = plt.subplots(figsize=(8, 7))
im = ax.imshow(confusion_percent, cmap="Blues", vmin=0.0, vmax=1.0)
ax.set_xticks(np.arange(NUM_CLASSES), CLASS_NAMES, rotation=45, ha="right")
ax.set_yticks(np.arange(NUM_CLASSES), CLASS_NAMES)
ax.set_xlabel("Predicted label")
ax.set_ylabel("True label")
ax.set_title("Fashion-MNIST confusion matrix")
fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="fraction of true class")

for i in range(NUM_CLASSES):
    for j in range(NUM_CLASSES):
        value = confusion_percent[i, j]
        if value >= 0.08 or i == j:
            ax.text(j, i, f"{100 * value:.0f}%", ha="center", va="center", fontsize=8, color="white" if value > 0.45 else "black")

fig.tight_layout()
plt.show()

You should see a dark diagonal with the off-diagonal cells mostly pale, and the darkest off-diagonal cells clustered among shirt, T-shirt/top, pullover, and coat.

8. Compare float32 and bfloat16 compute

float32 is the default for beginner training loops. bfloat16 uses fewer bits, so it can reduce memory traffic and may use faster hardware paths on supported NVIDIA GPUs. It is not automatically faster for every model, especially tiny ones, so in general: change the dtype, warm up, measure.

In this simple version, parameters stay in float32. During the forward pass, weights and activations are cast to the selected compute dtype, and logits are cast back to float32 before the loss. Many larger training systems use this same idea with more careful handling of numerically sensitive operations.

@partial(jax.jit, static_argnames=("compute_dtype",))
def train_step_mixed(params, opt_state, batch, compute_dtype=jnp.float32):
    """Same as `train_step`, but `compute_dtype` is a static argument."""
    (loss, metrics), grads = jax.value_and_grad(loss_with_metrics, has_aux=True)(
        params, batch, compute_dtype=compute_dtype
    )
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    metrics = {
        "loss": loss,
        "accuracy": metrics["accuracy"],
        "grad_norm": optax.global_norm(grads),
    }
    return params, opt_state, metrics


def time_mixed_precision(compute_dtype, steps=100):
    """Fresh init + one warmup compile for this dtype, then time `steps` steps and report examples/sec."""
    params_mp = init_mlp_params(jax.random.key(10))
    params_mp = jax.device_put(params_mp, device)
    opt_state_mp = optimizer.init(params_mp)

    params_mp, opt_state_mp, metrics = train_step_mixed(
        params_mp, opt_state_mp, first_batch, compute_dtype=compute_dtype
    )
    block_tree((params_mp, opt_state_mp, metrics))

    start = time.perf_counter()
    for step in range(steps):
        params_mp, opt_state_mp, metrics = train_step_mixed(
            params_mp, opt_state_mp, batch_at(step), compute_dtype=compute_dtype
        )
    params_mp, opt_state_mp, metrics = block_tree((params_mp, opt_state_mp, metrics))
    elapsed = time.perf_counter() - start
    return {
        "dtype": str(jnp.dtype(compute_dtype)),
        "loss": float(metrics["loss"]),
        "accuracy": float(metrics["accuracy"]),
        "elapsed": elapsed,
        "examples_per_sec": steps * BATCH_SIZE / elapsed,
    }


MIXED_PRECISION_STEPS = 100
mp_results = [
    time_mixed_precision(jnp.float32, steps=MIXED_PRECISION_STEPS),
    time_mixed_precision(jnp.bfloat16, steps=MIXED_PRECISION_STEPS),
]

show_table(
    ["Compute dtype", "Final loss", "Accuracy", "Elapsed seconds", "Examples/sec"],
    [
        (
            r["dtype"],
            f"{r['loss']:.4f}",
            f"{100 * r['accuracy']:.1f}%",
            f"{r['elapsed']:.3f}",
            f"{r['examples_per_sec']:,.0f}",
        )
        for r in mp_results
    ],
    title="Mixed-precision timing after warmup",
    aligns=["left", "right", "right", "right", "right"],
)
show_bars([(r["dtype"], r["examples_per_sec"]) for r in mp_results], "Mixed-precision examples/sec", "examples/s")

Compare the two rows. Both should reach a similar loss and accuracy, since the parameters and the loss stay in float32 either way.

Check what stayed on the GPU

The training loop returned new parameter and optimizer-state PyTrees. They are still JAX arrays on the GPU. The metrics only become Python values when you explicitly log them.

In real input pipelines, batches often start on the host. That is fine, but transfer a whole batch at a time and avoid conversions such as np.asarray(loss) or float(loss) inside the hot part of the loop.

def devices_in_tree(tree):
    """Set of devices that any JAX-array leaf in `tree` currently lives on."""
    devices = set()
    for leaf in jax.tree_util.tree_leaves(tree):
        if hasattr(leaf, "devices"):
            devices.update(leaf.devices())
    return devices


show_table(
    ["Object", "Where its arrays live"],
    [
        ("trained params", devices_in_tree(params_train)),
        ("optimizer state", devices_in_tree(opt_state)),
        ("training batches", x_train_batches.devices()),
        ("test batches", x_test_batches.devices()),
    ],
    title="Device placement check",
)

print(f"final training loss = {float(final_metrics['loss']):.4f}")

Every row of the placement table should name a CUDA device. Nothing quietly moved to the host during training, which is the property you want before you scale this loop up in the later codelabs.

9. Clean up

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.

10. Congratulations

You moved from individual JAX concepts to a complete GPU training loop on a real dataset, and you trained an MLP on Fashion-MNIST end to end.

What you've learned

• How to store model parameters as a PyTree of JAX arrays and keep fixed-size batches on the GPU
• How to write a scalar loss function for the objective you want to optimize
• When to use jax.grad (gradients only) and when to use jax.value_and_grad (loss and gradients from one pass)
• How to put an Optax AdamW update inside a single jax.jit-compiled training step
• How to measure throughput honestly: warm up first, use block_until_ready() before stopping the clock, and gate host reads behind a logging condition
• How examples/sec relates to tokens/sec, and why the tokens/sec number here is a projection rather than a measurement
• How to visualize predictions and a confusion matrix so the metrics connect back to actual examples
• Why bfloat16 is a useful option to measure rather than a guaranteed speedup, and why the parameters stay in float32

Next steps

• Codelab 5: Speed up attention on GPU with cuDNN and TransformerEngine where you will take the same compiled-step discipline into the attention kernels that dominate transformer training
• Change BATCH_SIZE and re-run. Larger batches often improve GPU utilization until memory becomes the limit, and each new shape costs one recompilation
• Change HIDDEN1 and HIDDEN2. More matrix multiplication usually makes the GPU busier, so watch what happens to examples/sec
• Change log_every in train_many_steps. More logging means more host synchronization, and the throughput number should show it

Reference docs

• Fashion-MNIST repository
• JAX automatic differentiation
• JAX asynchronous dispatch
• JAX default dtypes and the X64 flag
• Optax getting started