1. Introduction
In "Run your first JAX program on NVIDIA GPUs with GKE" codelab, you wrapped a function in jax.jit on the GPU and saw the first call take far longer than every call after it. That is not an accident: JAX traces your Python function with abstract placeholders, hands the recorded program to XLA, and caches the compiled executable it runs on the GPU. In this codelab you open up that process, learn what puts an entry in the compile cache, and fix the two things that cost JAX users the most time such as accidental recompiles, and Python control flow on traced values.
What you'll do
- Watch tracing happen by putting a Python
printinside a jitted function - Measure compilation cost against cached execution cost on the GPU
- Identify what belongs to the compile cache key and what triggers a recompile
- Replace Python control flow on traced values with
jnp.whereandjax.lax.cond - Keep shapes stable with
jax.lax.scan, padding and masking, andstatic_argnums - Inspect what JAX traced with
jax.make_jaxpr
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 Codelab 1: Run your first JAX program on NVIDIA GPUs with GKE, or an equivalent JAX GPU environment
Estimated time to complete: 50 minutes.
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>
This codelab runs in the same environment as Codelab 1: Run your first JAX program on NVIDIA GPUs with GKE. If your GKE cluster and JupyterLab Pod are still running, skip ahead to Set up and verify the GPU. Otherwise, provision the environment now.
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://, 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.
Set up and verify the GPU
Import JAX, NumPy, and a few standard-library helpers, then confirm you are on a GPU.
import time
from functools import partial
import jax
import jax.numpy as jnp
import numpy as np
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"Devices: {devices}")
assert gpu_devices, f"This lab assumes a GPU backend. Available devices: {devices}"
print(f"GPU devices: {gpu_devices}")
You should see gpu as the default backend and at least one CudaDevice in the device list. This codelab only needs one GPU, so it is fine if the node exposes more.
3. Watch jax.jit trace your function
When you call a plain, non-jitted JAX function, each operation runs through Python and dispatches to the GPU as it executes. jax.jit changes that. Instead of running your function with real arrays, it traces the function: JAX calls it once with abstract placeholders that carry only a shape and a dtype, and records every JAX operation you perform on those placeholders into an intermediate representation called a jaxpr.
JAX lowers the jaxpr to StableHLO, hands that lowered program to XLA, and XLA compiles an optimized executable for the target device. XLA may fuse operations, but a compiled function can still lower to multiple GPU kernels. From then on, calling the function jumps straight to the cached executable.
Every JIT call therefore has three phases:
Phase | What it does | What happens |
Trace | Record the computation | Python runs once and JAX records every operation on abstract placeholders into a |
Compile | Lower to a GPU executable | JAX lowers the |
Execute | Reuse the cached executable | Every later call with matching shapes and dtypes skips trace and compile, and runs the cached program. |
Trace is why a Python print inside a JIT function only fires on the first call. Trace and Compile together are why the first call is slow. Execute is why every call after it is fast.
See tracing in action
Prove to yourself that the function body only runs once per input signature. Put a Python-level print inside the function: it executes during tracing, but it is not part of the compiled GPU program, so later calls with the same shape and dtype print nothing.
@jax.jit
def f(x):
"""Jitted demo function that prints during tracing so we can see exactly when JAX retraces."""
# This print runs during tracing only not on every GPU execution.
print(f" tracing with shape={x.shape} dtype={x.dtype}")
return x ** 2 + 1
print("Call 1 (new shape):")
_ = f(jnp.arange(4, dtype=jnp.float32)).block_until_ready()
print("Call 2 (same shape):")
_ = f(jnp.arange(4, dtype=jnp.float32)).block_until_ready()
print("Call 3 (new shape):")
_ = f(jnp.arange(5, dtype=jnp.float32)).block_until_ready()
You should see output similar to:
Call 1 (new shape): tracing with shape=(4,) dtype=float32 Call 2 (same shape): Call 3 (new shape): tracing with shape=(5,) dtype=float32
The print fires on call 1, the first time JAX sees shape (4,) with dtype float32, and on call 3, the first time it sees shape (5,). On call 2, JAX finds an existing compiled executable and skips both tracing and compilation.
4. Measure compilation against cached execution
The new signature cost is real, and it is where JAX slow reports come from. Measure how much of the first call is compilation and how much is execution.
The function below chains 20 nonlinearities so that compilation is visibly more expensive than execution.
def heavy(x):
"""20 chained nonlinearities so the first-call compilation is visibly more expensive than the cached execution."""
y = x
for _ in range(20):
y = jnp.sin(y) * jnp.cos(y) + jnp.tanh(y)
return y
heavy_jit = jax.jit(heavy)
x = jnp.arange(1_000_000, dtype=jnp.float32)
# Empty in-process cache so a re-run shows the first-call compile cost again.
jax.clear_caches()
t0 = time.perf_counter()
_ = heavy_jit(x).block_until_ready()
first_ms = (time.perf_counter() - t0) * 1000
t0 = time.perf_counter()
for _ in range(20):
_ = heavy_jit(x).block_until_ready()
cached_ms = (time.perf_counter() - t0) * 1000 / 20
print(f"First call (compile + execute): {first_ms:8.2f} ms")
print(f"Cached call (execute only): {cached_ms:8.2f} ms")
print(f"Compilation cost (approx): {first_ms - cached_ms:8.2f} ms")
You should see a first-call time that is much larger than the cached-call time. The gap between them is roughly what XLA spent compiling.
For tiny functions that gap is a few tens of milliseconds; for a full transformer training step it can easily be several seconds. The good news is that you pay it once per shape and dtype combination, not once per call. The rest of this codelab is about not paying it more often than you have to.
5. Find out what triggers a recompile
JAX keys the compile cache on a structural signature of the inputs: their shapes, their dtypes, and any arguments marked as static. If the signature matches one JAX has seen before, the cached executable runs. If anything changes, JAX traces and compiles again.
Three things commonly trigger a recompile:
Cache key part | What changes | Effect |
Shape | Different shape |
|
Dtype | Different dtype |
|
Static arg | Different static value | The value of any |
The values of regular array inputs do not matter. Two (32, 128) float32 arrays with completely different contents hit the same compiled executable.
Watch a recompile happen. The loop below calls one jitted function with five arrays, three of them a shape JAX has not seen yet.
@jax.jit
def f(x):
"""Simple jitted scalar function used to demonstrate one compile per new input shape (a new dtype would trigger the same recompile)."""
return jnp.sum(x ** 2)
# clear JAX's in-process compilation cache.
jax.clear_caches()
# Feed in a few different shapes and measure each call.
shapes = [(100,), (200,), (100,), (200,), (300,)]
for s in shapes:
x = jnp.ones(s, dtype=jnp.float32)
t0 = time.perf_counter()
_ = f(x).block_until_ready()
dt = (time.perf_counter() - t0) * 1000
print(f"shape={str(s):8s} {dt:7.2f} ms")
You should see three slow calls, one per new shape, and two fast ones, for the repeated (100,) and (200,).
Real workloads do this accidentally all the time: variable-length sequences, the last batch in an epoch, ragged tokenization output. The fix in almost every case is don't let the shape change.
6. Replace Python control flow on traced values
During tracing, your function's inputs are not concrete arrays. They are abstract values with a known shape and dtype. Any Python construct that needs to compare those contents numerically (if, while, bool(x), int(x)) breaks tracing.
Here is what that looks like. This ReLU is wrong by design:
@jax.jit
def relu_bad(x):
"""ReLU using a Python `if` on a traced value with JIT errors out at trace time."""
if x > 0:
return x
return jnp.zeros_like(x)
try:
print(relu_bad(jnp.array(1.0)))
except Exception as e:
print(f"{type(e).__name__}: {str(e).splitlines()[0]}")
You should see a TracerBoolConversionError. The message points at the if: JAX cannot decide which branch to keep when the value is abstract.
Express the choice as data with jnp.where
The fix is to express the choice as data, not as Python control flow. For a small elementwise selection like ReLU, jnp.where is the cleanest tool. Both branches always run, and the predicate tells JAX which one to use at each position.
@jax.jit
def relu(x):
"""ReLU using `jnp.where` with both branches are computed so tracing works."""
return jnp.where(x > 0, x, 0.0)
print(relu(jnp.array([-1.0, -0.5, 0.0, 0.5, 1.0])))
This time there is no error. The two negative entries and the zero come back as 0., and 0.5 and 1.0 pass through unchanged.
Pick a real branch with jax.lax.cond
For branches that compute very different things, where running both would be wasteful, use jax.lax.cond. Both branch functions are traced, but at runtime lax.cond represents an XLA conditional, so normally only the selected branch executes. One caveat: under vmap, cond may be converted to a select-like operation rather than a real branch.
@jax.jit
def soft_or_sharp(x, sharp):
"""Switch between hard ReLU and softplus inside the compiled graph via `lax.cond`, controlled by a traced bool."""
# `sharp` is a scalar bool and lax.cond compiles to a real if-then-else
return jax.lax.cond(
sharp,
lambda x: jnp.where(x > 0, x, 0.0),
lambda x: jax.nn.softplus(x),
x,
)
x = jnp.array([-1.0, 0.5, 2.0])
print(f"sharp=True: {soft_or_sharp(x, jnp.array(True))}")
print(f"sharp=False: {soft_or_sharp(x, jnp.array(False))}")
You should see two different arrays: the sharp=True line clamps the negative entry to zero, and the sharp=False line returns small positive softplus values everywhere.
7. Keep long loops compact with lax.scan
For loops over traced data, use structured control-flow primitives such as jax.lax.while_loop, jax.lax.fori_loop, and jax.lax.scan.
In fact, a Python for loop with a static bound is valid inside jit, but JAX unrolls the loop while tracing. That means 200 loop iterations become roughly 200 repeated blocks in the compiled program. lax.scan keeps the loop as a loop-like primitive, which usually compiles much faster for long fixed-length loops.
# python_for_loop compile time scales with NUM_STEPS while scan_loop compile time stays roughly constant. Try NUM_STEPS = 2000 to see the gap widen.
NUM_STEPS = 200
@jax.jit
def python_for_loop(x):
"""Python `for` loop inside jit."""
y = x
for _ in range(NUM_STEPS):
y = jnp.sin(y) + 0.01 * y
return y
@jax.jit
def scan_loop(x):
"""Same logic expressed with `lax.scan`."""
def body(y, _):
y = jnp.sin(y) + 0.01 * y
return y, None
y, _ = jax.lax.scan(body, x, xs=None, length=NUM_STEPS)
return y
x = jnp.ones((1024,), dtype=jnp.float32)
jax.clear_caches()
t0 = time.perf_counter()
_ = python_for_loop(x).block_until_ready()
python_for_ms = (time.perf_counter() - t0) * 1000
jax.clear_caches()
t0 = time.perf_counter()
_ = scan_loop(x).block_until_ready()
scan_ms = (time.perf_counter() - t0) * 1000
print(f"Python for loop first call: {python_for_ms:8.2f} ms")
print(f"lax.scan first call: {scan_ms:8.2f} ms")
Both numbers include compilation, and both functions compute the same recurrence. The unrolled Python loop has to compile a much larger program, so its first call is the slower of the two.
The important difference is compile-time structure: the Python loop is unrolled during tracing, while lax.scan lowers to a loop primitive. For short loops, a Python for loop is often fine. For long differentiable loops, lax.scan is usually the better default.
8. Stabilize shapes with padding and static arguments
Real workloads vary in shape. For example, the last batch in an epoch is smaller, sequences have different lengths. Every one of those triggers a fresh compile if you let the shape leak through to JAX. The standard fix is to pad inputs to a fixed shape and mask the unused positions.
MAX_LEN = 16
@jax.jit
def masked_mean(x, mask):
"""Mean of `x` ignoring positions where `mask==0`."""
# Always called with shape (MAX_LEN,) - no recompile when actual length varies
return jnp.sum(x * mask) / jnp.maximum(jnp.sum(mask), 1.0)
def pad(seq):
"""Right-pad a variable-length list of floats to `MAX_LEN` and return the padded array plus a 0/1 mask."""
actual_len = len(seq)
if actual_len > MAX_LEN:
raise ValueError(f"sequence length {actual_len} exceeds MAX_LEN={MAX_LEN}")
pad_len = MAX_LEN - actual_len
x = jnp.concatenate([
jnp.asarray(seq, dtype=jnp.float32),
jnp.zeros(pad_len, dtype=jnp.float32),
])
mask = jnp.concatenate([
jnp.ones(actual_len, dtype=jnp.float32),
jnp.zeros(pad_len, dtype=jnp.float32),
])
return x, mask
# Several different sequence lengths, but a single compiled function handles them all
for seq in [[1.0, 2.0, 3.0], [10.0] * 8, [5.0, -2.0]]:
x, mask = pad(seq)
print(f"len={len(seq):2d} mean={masked_mean(x, mask):.3f}")
You should see one line per sequence, each reporting the mean of the real values only — the padding does not drag the average toward zero.
All three calls hit the same compiled executable because the shape on the device is always (MAX_LEN,). Only the mask changes. This pattern shows up at every scale, from a 16-element simple mean here to padded attention masks in large-scale transformer training.
When you want a recompile: static_argnums
Padding is how you avoid a recompile you never asked for but static_argnums is how you ask for one on purpose.
Sometimes a parameter genuinely is a Python-side constant like a layer count, a precision flag, or a kernel size and you want JAX to bake its value into the compiled program. Mark those arguments with static_argnums, or static_argnames for keyword arguments. JAX hashes the value of those arguments into the cache key, so each distinct value gets its own compiled executable.
@partial(jax.jit, static_argnums=0)
def power(n: int, x):
"""Repeated squaring with `n` is static so JAX unrolls the loop and compiles a fresh program per value of `n`."""
# `n` is a Python int and JAX bakes it into the trace and unrolls the loop
y = x
for _ in range(n):
y = y * y
return y
jax.clear_caches()
for n in (2, 3, 2): # n=2 reuses the cache the second time
t0 = time.perf_counter()
_ = power(n, jnp.arange(4, dtype=jnp.float32)).block_until_ready()
print(f"n={n}: {(time.perf_counter() - t0) * 1000:7.2f} ms")
You should see the first two calls compile and the third one, which repeats n=2, come back fast.
Each new value of n triggers a compile, but for fixed configurations that is exactly what you want: the loop unrolls completely and XLA can see every operation. The trade-off is straightforward: do not put a continuously varying value in static_argnums, or you will recompile on every call.
9. Inspect the trace with jax.make_jaxpr
When something compiles differently than you expect, jax.make_jaxpr lets you see the trace before XLA touches it. A jaxpr is a JAX-level compiler intermediate representation: a typed, functional representation of what JAX staged out before lowering to StableHLO and then XLA.
It is not the final optimized GPU code, but it is very useful for understanding what JAX traced.
def f(x):
return jnp.tanh(x) * jnp.sin(x) + jnp.log1p(x * x)
print(jax.make_jaxpr(f)(jnp.arange(4, dtype=jnp.float32)))
You should see a small typed program: one primitive per line — tanh, sin, the multiplications, log1p, and the final add — each annotated with the array type it produces.
If you ever suspect that JAX is recompiling because a shape or dtype changed unexpectedly, comparing two jaxprs from a "fast" and a "slow" call usually pinpoints the culprit. The same trick works for figuring out why a transformation such as grad or vmap is producing more work than you intended.
10. 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.
11. Congratulations
You built the mental model behind jax.jit: JAX traces your Python function, turns the traced computation into compiler input, and caches the compiled executable for matching input signatures.
What you've learned
- How to trace a Python function with
jax.jit, and why Python side effects such asprintrun during tracing rather than on every execution - How to distinguish compilation time from cached execution time with simple
block_until_ready()timing - What the compile cache keys on: input PyTree structure, shapes, dtypes, and static argument values
- How to avoid traced-value control-flow errors by replacing Python branching with
jnp.whereorjax.lax.cond, and whyjnp.wherecan leak NaNs into a gradient - How
lax.scankeeps a long fixed-length loop compact instead of unrolling it into the compiled program - How to stabilize input shapes with padding and masking, and how to bake Python-side constants into the trace with
static_argnumswhen you intentionally want a separate executable - How to inspect the JAX-level trace with
jax.make_jaxprwhen a function compiles differently than you expect
Next steps
- In Codelab 3: Profile and debug JAX on GPU with XProf and Nsight Systems you will learn to watch tracing, compilation, and kernel execution show up in a real profile
- Raise
NUM_STEPSfrom 200 to 2000 and re-run the loop comparison, to widen the gap betweenlax.scanand the unrolled Python loop - Put a continuously varying value in
static_argnums: callpowerwithntaken from a counter that increments every call, and watch every single call recompile