Core Foundation · Local CPULaptop CPULesson 5 of 5 in Phase 02

Phase 02 / JAX transformations

Tracing, static arguments, and recompilation

Revision v0.7.0 · 2026-10-20

Lesson · about 60 minutes · CPU practice

Explain the staging boundary and repair configuration versus runtime-control-flow mistakes.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: Compile a function with jit
Hardware: CPU

Download notebook

Other downloads

Python worked examples Lesson text

Setup help: Python, notebook, or Cloud TPU VM

Use the tested CPU packages (the recorded environment uses Python 3.14.3). Other Python versions and accelerator environments require separate checks. The optional Cloud TPU specialization explains the target-specific execution boundaries. Logical CPU devices practice partitioning; they do not simulate TPU performance.

Download the repository ZIP or use your existing checkout. Extract it and open a terminal in the folder containing requirements-cpu.txt. The beginner ZIP is separate; these lessons use the full repository folder.

macOS / Linux terminal · repository folder

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt

PowerShell · repository folder

py -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txt

After activation, run the complete worked example from that same folder:

macOS / Linux terminal · activated environment

python phases/02-transforms/05-tracing-static-arguments-and-recompilation/code/main.py

PowerShell command · repository folder

.\.venv\Scripts\python.exe phases/02-transforms/05-tracing-static-arguments-and-recompilation/code/main.py

For an incremental build, create main.py in your own practice folder, append each build block in order, and run python main.py with this environment active. Keep that file separate from the downloaded reference solution.

Download this lesson’s notebook. Open it with a Python kernel using the course environment. Run cells from the top after restarting the kernel; experiments depend on earlier definitions.

For local Jupyter, first install the course requirements using the Local Python steps, then use your environment’s Python to start Jupyter:

macOS / Linux terminal · activated course environment

python -m pip install jupyterlab
python -m jupyter lab

PowerShell command · repository folder

.\.venv\Scripts\python.exe -m pip install jupyterlab
.\.venv\Scripts\python.exe -m jupyter lab

Select the downloaded transforms-05.ipynb file in JupyterLab. The notebook contains reference solutions; make your prediction and attempt each practice task before opening them. Restart the kernel before changing logical device setup.

Check: Sum: $6.0$ ; Mean: $2.0$ . These mode choices can use separate compiled specializations.

Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.

Phase 02.05 Stage 1 of 5 · Understand

1 Understand 2 Build & Run 3 Experiment 4 Exercise & Check 5 Evidence

What you will be able to do

Distinguish runtime array values from shape/dtype metadata and static configuration.
Inspect a staged computation with make_jaxpr.
Explain a Python value-dependent branch failure under jit.
Use static configuration or JAX control flow according to the kind of decision.

The problem

A function works until you put a Python if statement inside jit. Why does that change things? We’ll separate what Python knows while tracing from what the program learns at runtime. Then you’ll choose between a small static configuration and an array-based decision, rather than making every difficult argument static.

The idea

Tracing observes how Python builds an array computation from abstract inputs. Static arguments are part of that program-building decision; dynamic array values are data for the resulting computation. Keep a small call ledger to separate these roles.

Build a call ledger before diagnosing recompilation

Record each call's shape, dtype and static choices next to the observed trace count. Change only values first. Then change one shape, then one static option. This controlled sequence makes a trace increase interpretable; changing everything at once does not.

A Python branch needs a concrete decision during tracing. If the decision is intended to vary with array data at execution time, use an appropriate JAX control-flow operation instead of converting a tracer to a Python boolean. If it genuinely selects a small number of program variants, an explicit static argument may be suitable.

Do not turn every input into a static value to silence an error. Many distinct static values can create many specializations. Trace counts are useful observations, but they are not a universal count of all compilation work or a cross-process cache guarantee.

\theta_{t+1} = \theta_t - \eta \cdot \frac{1}{B}\sum_{i=1}^{B} \nabla_\theta \ell(\theta_t, x_i, y_i)

Separate metadata, configuration, and data

A tensor shape such as $(3,)$ describes the size of the computation; its values $[1, 2, 3]$ are runtime data. A string mode such as sum describes which computation to construct. In this lesson mode is a small finite configuration, so marking it static is reasonable.

If mode changes rarely while batches change constantly, static configuration lets a stable compiled specialization accept many batches. If you make frequently changing data static, you encourage distinct specializations and may run into hashability constraints. A JAX array is not a convenient static argument just because a Python branch wants to inspect it.

Read a tiny staged program

make_jaxpr exposes a representation of the traced array operations. A Python print inside the function may run while tracing, but it does not become an ordinary numerical operation in that representation. A multiply and a reduce_sum should appear in a multiply-then-sum example.

The exact printed formatting and primitive names are not a stable teaching contract. Read it for inputs, outputs, and the broad sequence of numerical operations. The point is to observe that the computation is staged, not to memorize an internal dump or treat the representation as a complete performance explanation.

Why a Python if can fail

Consider if $x$ > $0$ where $x$ is a dynamic scalar argument. Python wants one Boolean immediately to choose which branch to execute while tracing. The tracer represents a runtime value, so that truth test cannot generally be answered at trace time. The resulting conversion error is evidence of a staging mismatch, not evidence that comparisons cannot run on an accelerator.

For a scalar runtime decision, lax.cond expresses the conditional in the numerical program. Its two branches must return compatible structures and shapes/dtypes. Both branch functions are traced; for an ordinary scalar condition the selected branch is the runtime computation. Under batching, a conditional can be transformed into selection behavior, so do not infer cost or side-effect behavior from the Python-like notation alone.

Make configuration choices explicit

The worked reduce_values function chooses sum or mean using a static mode string. Each mode can have its own specialization. It should reject unknown modes instead of silently treating any typo as mean. This is a small input-contract improvement that prevents plausible-looking wrong results.

Keep function identity stable, configure static_argnames at the transformation boundary, and pass mode deliberately. Shape and dtype changes can also require specializations. Different element values with compatible abstract input properties ordinarily do not, though caches and transformation details make exact global compilation counts a poor public guarantee.

Inspect tracing without confusing it with runtime logging

A Python print in a jitted function is a trace-time observation. If you call the function again with compatible inputs, its absence is consistent with reuse. It is not a reliable counter of all device executions or all compilations across the application. JAX provides debug facilities for observing runtime values.

Use trace observations for a small diagnostic experiment and rely on supported compiler logging/profiling when investigating a real training job. Do not make application behavior depend on a side effect inside a transformed numerical function.

Carry this boundary into the next phase

The upcoming state lessons will make randomness, parameter structure, loops, and branches explicit. The same principle applies: information needed at runtime belongs in the numerical state or arguments, while small program configuration can belong in specialization.

Before reaching for a workaround, classify the value and the question: is Python constructing the computation, or is the computation choosing a result based on runtime data? That classification leads to an implementation you can explain and test.

Compiled branch outputs: sum mode vs. mean mode

Predict before looking: Predict the ratio between compiled(x, 'sum') and compiled(x, 'mean') for the input array x.

Read the figure

Static mode 'sum' returns the sum and 'mean' returns the mean without Python control-flow tracing errors.

Connect it to the computation

Marking mode as a static argument in jax.jit allows each distinct mode string to specialize and compile its own branch cleanly.

Open full-size figure Download data and provenance

The notebook contains the data experiment, plotting code, and recorded figure output. Rerun it to test your changes.

Where to put & run this code: Work in phases/02-transforms/05-tracing-static-arguments-and-recompilation/code/main.py (or downloaded exercises/transforms-05.py / exercises/transforms-05.ipynb) and execute from the repository root with python3 phases/02-transforms/05-tracing-static-arguments-and-recompilation/code/main.py.

1 Step 1: Set up imports and input tensors

Import the required JAX modules and define the initial inputs for tracing, static arguments, and recompilation.

main.py · append build step 1

import jax
import jax.numpy as jnp

Establishing explicit input shapes and dtypes first makes the downstream transformation contract deterministic.

2 Step 2: Apply the core JAX transformation

Write the core computation and transformation step over the initialized inputs.

main.py · append build step 2

def reduce_values(x, mode):
    if mode == "sum":
        return x.sum()
    if mode == "mean":
        return x.mean()
    raise ValueError("mode must be sum or mean")
# Wrap with `jax.jit` (`compiled`) so XLA traces and compiles the function.
compiled = jax.jit(reduce_values, static_argnames=("mode",))

Code building blocks in this step

x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
jax.jit(fn) / @jax.jit — Traces fn with abstract shapes and compiles a fused XLA executable cached by input shape and dtype.

This stage executes the primary numerical transformation and binds the intermediate outputs.

3 Step 3: Verify shapes and numerical invariants

Check that the resulting arrays satisfy the expected shape, dtype, and numerical tolerances.

main.py · append build step 3

x = jnp.array([1., 2., 3.])
assert jnp.allclose(compiled(x, mode="sum"), 6.)
assert jnp.allclose(compiled(x, mode="mean"), 2.)

Code building blocks in this step

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.

These assertions lock in the exact numerical contract before you run the full experiment and variations.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build

import jax
import jax.numpy as jnp

def reduce_values(x, mode):
    if mode == "sum":
        return x.sum()
    if mode == "mean":
        return x.mean()
    raise ValueError("mode must be sum or mean")
# Wrap with `jax.jit` (`compiled`) so XLA traces and compiles the function.
compiled = jax.jit(reduce_values, static_argnames=("mode",))

x = jnp.array([1., 2., 3.])
assert jnp.allclose(compiled(x, mode="sum"), 6.)
assert jnp.allclose(compiled(x, mode="mean"), 2.)

Run the example

python3 phases/02-transforms/05-tracing-static-arguments-and-recompilation/code/main.py

About this command

This shortcut uses an activated macOS/Linux course environment in the repository folder. For Windows, use the explicit interpreter command in “Setup help” above. The script includes worked examples and reference solutions.

Inspect the complete assembled example

main.py · complete Python example

# Tracing, static arguments, and recompilation: Tracing observes how Python builds an array computation from...
import jax
import jax.numpy as jnp
def reduce_values(x, mode):
    if mode == "sum":
        return x.sum()
    if mode == "mean":
        return x.mean()
    raise ValueError("mode must be sum or mean")
# Wrap with `jax.jit` (`compiled`) so XLA traces and compiles the function.
compiled = jax.jit(reduce_values, static_argnames=("mode",))
x = jnp.array([1., 2., 3.])
print("Sum:", float(compiled(x, mode="sum")))
print("Mean:", float(compiled(x, mode="mean")))
assert jnp.allclose(compiled(x, mode="sum"), 6.)
assert jnp.allclose(compiled(x, mode="mean"), 2.)

Expected result: Sum: $6.0$ ; Mean: $2.0$ . These mode choices can use separate compiled specializations.

Recorded reference execution

Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.

2026-10-20T10:45:24.879206+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

Sum: 6.0
Mean: 2.0
Tracing sum_squares; shape: (3,)
{ lambda ; a:f32[3]. let
    b:f32[3] = mul a a
    c:f32[] = reduce_sum[axes=(0,) out_sharding=None] b
  in (c,) }
Expected dynamic Boolean conversion error
PASS: transforms-05

Download execution record · The notebook includes executed cells and figure outputs. Compare your run with these values; timings depend on the machine.

Look at the staged numerical work

Look at the staged numerical work · Python experiment

# Experiment — Look at the staged numerical work: The second call supplies new values with the same shape/dtype.
def sum_squares(values):
    print("Tracing sum_squares; shape:", values.shape)
    return jnp.sum(values * values)
print(jax.make_jaxpr(sum_squares)(x))
# Wrap with `jax.jit` (`staged_squares`) so XLA traces and compiles the function.
staged_squares = jax.jit(sum_squares)
assert jnp.allclose(staged_squares(x), 14.)
assert jnp.allclose(staged_squares(x + 1.), 29.)

Expected: The representation contains numerical multiplication and reduction. Trace prints can appear during staging, not as one guaranteed print per call.

The second call supplies new values with the same shape/dtype. The numerical result changes while the computation structure can be reused.

Reproduce and repair a dynamic branch

Reproduce and repair a dynamic branch · Python experiment

# Experiment — Reproduce and repair a dynamic branch: The test expects a particular category of error and confirms the...
def python_branch(value):
    if value > 0.:
        return value ** 2
    return -value
assert python_branch(-2.) == 2.
# Run the boundary check and catch the expected exception:
try:
    jax.jit(python_branch)(jnp.array(-2.))
except jax.errors.TracerBoolConversionError:
    print("Expected dynamic Boolean conversion error")
else:
    raise AssertionError("Expected a traced Python branch to fail")
def runtime_branch(value):
    return jax.lax.cond(value > 0., lambda z: z ** 2, lambda z: -z, value)
# Wrap with `jax.jit` (`compiled_branch`) so XLA traces and compiles the function.
compiled_branch = jax.jit(runtime_branch)
assert jnp.allclose(compiled_branch(jnp.array(-2.)), 2.)
assert jnp.allclose(compiled_branch(jnp.array(3.)), 9.)

Expected: The Python branch fails under jit with a tracer Boolean conversion error. The staged conditional returns $2$ and $9$ .

The test expects a particular category of error and confirms the repair for both branches. A broad catch-all would hide unrelated failures.

Hands-On Exercise · Modify & Verify Run in the declared lesson environment

Make it yours

Where to put your code: Append your solution at the bottom of phases/02-transforms/05-tracing-static-arguments-and-recompilation/code/main.py (or inside the Make it yours cell of exercises/transforms-05.ipynb) and run python3 phases/02-transforms/05-tracing-static-arguments-and-recompilation/code/main.py to verify your prediction before opening the reference solution.

Foundation · Extend reduce_values with a static max mode while preserving sum and mean. Reject unknown modes. Verify all three modes on $[1, 2, 3]$ and on $[-4, -1, -2]$ .

Show a writing guide and starter scaffold

How to write this exercise code: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
jax.jit(fn) / @jax.jit — Traces fn with abstract shapes and compiles a fused XLA executable cached by input shape and dtype.

2. Step-by-step implementation plan:

Wrap with jax.jit (choose) so XLA traces and compiles the function.
Run the boundary check and catch the expected exception:

Starter code scaffold · copy into your file & fill in the TODOs

# Exercise solution: Foundation · Extend reduce_values with a static max mode while...
def choose_reduction(values, mode):
    if mode = ...  # TODO: compute if mode
        return ...  # TODO: return computed result
    if mode = ...  # TODO: compute if mode
        return ...  # TODO: return computed result
    if mode = ...  # TODO: compute if mode
        return ...  # TODO: return computed result
    raise ValueError("unknown mode")
# Wrap with `jax.jit` (`choose`) so XLA traces and compiles the function.
choose = jax.jit(...)  # TODO: compute choose
for values in (x, jnp.array([-4., -1., -2.])):
    for mode, reference in (("sum", jnp.sum), ("mean", jnp.mean), ("max", jnp.max)):
        assert jnp.allclose(choose(values, mode=mode), reference(values))  # TODO: complete assertion check
# Run the boundary check and catch the expected exception:
try:
    choose(x, mode = ...  # TODO: compute choose(x, mode
except ValueError:
    pass
else:
    raise AssertionError("Unknown mode should fail")

Reveal the reference solution

Reference solution

# Exercise solution: Foundation · Extend reduce_values with a static max mode while...
def choose_reduction(values, mode):
    if mode == "sum":
        return values.sum()
    if mode == "mean":
        return values.mean()
    if mode == "max":
        return values.max()
    raise ValueError("unknown mode")
# Wrap with `jax.jit` (`choose`) so XLA traces and compiles the function.
choose = jax.jit(choose_reduction, static_argnames=("mode",))
for values in (x, jnp.array([-4., -1., -2.])):
    for mode, reference in (("sum", jnp.sum), ("mean", jnp.mean), ("max", jnp.max)):
        assert jnp.allclose(choose(values, mode=mode), reference(values))
# Run the boundary check and catch the expected exception:
try:
    choose(x, mode="typo")
except ValueError:
    pass
else:
    raise AssertionError("Unknown mode should fail")

Keep runtime values dynamic

Practice

Implement an elementwise absolute value with a numerical selection, jit it, and test mixed-sign vectors. Explain why making the whole vector static is inappropriate.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
jax.jit(fn) / @jax.jit — Traces fn with abstract shapes and compiles a fused XLA executable cached by input shape and dtype.

2. Step-by-step implementation plan:

Wrap with jax.jit (compiled_absolute) so XLA traces and compiles the function.

Starter code scaffold · copy into your file & fill in the TODOs

# Keep runtime values dynamic (Practice): The vector values are runtime data.
def elementwise_absolute(values):
    return ...  # TODO: return computed result
# Wrap with `jax.jit` (`compiled_absolute`) so XLA traces and compiles the function.
compiled_absolute = jax.jit(...)  # TODO: compute compiled_absolute
for values in (jnp.array([-2., 0., 3.]), jnp.array([5., -4., -1.])):
    assert jnp.allclose(compiled_absolute(values), jnp.abs(values))  # TODO: complete assertion check

Show a hint

jnp.where selects elementwise; both candidate arrays are computed expressions. It is not a Python branch per element.

Reveal reference solution and reasoning

Keep runtime values dynamic · Python solution

# Keep runtime values dynamic (Practice): The vector values are runtime data.
def elementwise_absolute(values):
    return jnp.where(values >= 0., values, -values)
# Wrap with `jax.jit` (`compiled_absolute`) so XLA traces and compiles the function.
compiled_absolute = jax.jit(elementwise_absolute)
for values in (jnp.array([-2., 0., 3.]), jnp.array([5., -4., -1.])):
    assert jnp.allclose(compiled_absolute(values), jnp.abs(values))

The vector values are runtime data. Elementwise selection preserves a fixed output shape and does not need a static array.

Audit a shape-dependent specialization

Challenge

Run a jitted sum-of-squares function for lengths $3$ and $5$ , each with two different value sets. Compare to a loop calculation. Explain which changes can affect specialization and which are ordinary runtime data.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

jnp.arange(n, dtype=...) — Creates a 1-D JAX array of evenly spaced values [0, 1, ..., n-1] on the target device.
x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
jax.jit(fn) / @jax.jit — Traces fn with abstract shapes and compiles a fused XLA executable cached by input shape and dtype.

2. Step-by-step implementation plan:

Wrap with jax.jit (compiled_squares) so XLA traces and compiles the function.
Create evenly spaced index values in values.

Starter code scaffold · copy into your file & fill in the TODOs

# Audit a shape-dependent specialization (Challenge): Changing the length changes shape metadata and can need a...
def plain_sum_squares(values):
    return ...  # TODO: return computed result
# Wrap with `jax.jit` (`compiled_squares`) so XLA traces and compiles the function.
compiled_squares = jax.jit(...)  # TODO: compute compiled_squares
for length in (3, 5):
    for shift in (0., 1.):
        # Create evenly spaced index values in `values`.
        values = jnp.arange(...)  # TODO: compute values
        expected = sum(...)  # TODO: compute expected
        assert jnp.allclose(compiled_squares(values), expected)  # TODO: complete assertion check

Show a hint

Do not use elapsed time alone to infer whether tracing occurred. Values and shape are different inputs to your reasoning.

Reveal reference solution and reasoning

Audit a shape-dependent specialization · Python solution

# Audit a shape-dependent specialization (Challenge): Changing the length changes shape metadata and can need a...
def plain_sum_squares(values):
    return jnp.sum(values ** 2)
# Wrap with `jax.jit` (`compiled_squares`) so XLA traces and compiles the function.
compiled_squares = jax.jit(plain_sum_squares)
for length in (3, 5):
    for shift in (0., 1.):
        # Create evenly spaced index values in `values`.
        values = jnp.arange(length, dtype=jnp.float32) + shift
        expected = sum(float(v) ** 2 for v in values)
        assert jnp.allclose(compiled_squares(values), expected)

Changing the length changes shape metadata and can need a new specialization. Changing shift changes element values at the same shape/dtype. The test verifies correctness across both kinds of changes without inventing a guaranteed compilation count.

Diagnose the result

For a tracer Boolean conversion error, inspect the Python truth test. For an unhashable static argument, reconsider whether the object is really configuration. For repeated setup, inspect function recreation, shapes/dtypes, and static values. For missing Python logs, distinguish tracing from runtime execution. Always verify both branches and preserve unknown-mode errors rather than returning a plausible fallback.

Carry forward

Python executes while staging; numerical array operations form the runtime computation.
Static configuration and dynamic numerical data have different roles.
Use runtime control flow for runtime decisions and compatible branch returns.
Inspect evidence of tracing separately from numerical correctness and timing.

Keep your evidence

Download evidence template

Keep the small jaxpr observation, the expected tracer Boolean error, checks of both repaired branches, all static reduction modes including invalid-mode rejection, and shape/value-change cases. Do not infer compilation counts from timing alone.

Save your modified code, its output, and an explanation of your prediction. Checkpoint answers and exercise evidence are tracked separately.

Saved in this browser. Export or back up your portfolio anytime in My learning.

Save work for this lesson

Verified on local CPU. Stage to a Cloud TPU VM via Cloud TPU setup to record accelerator execution receipts (TODO: add automated Cloud TPU CI execution receipts).

Primary references

Phase 02 Integration Project

Bring this phase together: Compare scalar, vmap, autodiff, and changed-shape JIT behavior

Checkpoint to prove: Compose grad, vmap, and jit; explain what each transformation changes.

Work in projects/transforms-audit/starter/model.py. Complete stages 1 and 2 for independent calculations and changed-condition diagnosis, then retain the stage 3 evidence with --evidence-dir runs/transforms-audit. Stage 4 integrates the checks. Learner commands explicitly select the starter; --implementation solution is instructor reference execution. This bounded practice does not certify the entire phase or target hardware.

Open project workspace · stages 1, 2, 3, 4 Read project guide & figures

Tracing, static arguments, and recompilation

What you will be able to do

The problem

The idea

Build a call ledger before diagnosing recompilation

Separate metadata, configuration, and data

Read a tiny staged program

Why a Python if can fail

Make configuration choices explicit

Inspect tracing without confusing it with runtime logging

Carry this boundary into the next phase

Compiled branch outputs: sum mode vs. mean mode

Read the figure

Connect it to the computation

1 Step 1: Set up imports and input tensors

2 Step 2: Apply the core JAX transformation

3 Step 3: Verify shapes and numerical invariants

Run the example

Recorded reference execution

Look at the staged numerical work

Reproduce and repair a dynamic branch

Make it yours

Keep runtime values dynamic

Audit a shape-dependent specialization

Check your understanding

Diagnose the result

Carry forward

Keep your evidence

Primary references

Bring this phase together: Compare scalar, vmap, autodiff, and changed-shape JIT behavior