Training Systems · Local CPULaptop CPULesson 2 of 5 in Phase 08

Phase 08 / Performance diagnosis

Diagnose recompilation and host synchronization

Revision v0.7.0 · 2026-10-20

Lesson · about 80 minutes · CPU practice

Find a shape-changing call or host read and connect it to observed overhead.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: Benchmark asynchronous work correctly
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/08-performance/02-diagnose-recompilation-and-host-synchronization/code/main.py

PowerShell command · repository folder

.\.venv\Scripts\python.exe phases/08-performance/02-diagnose-recompilation-and-host-synchronization/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 performance-02.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: A JSON ledger of four calls and observed Python trace events. Checked scores include $140$ , $560$ and $506$ ; event deltas are recorded rather than fixed by the lesson.

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

Phase 08.02 Stage 1 of 5 · Understand

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

What you will be able to do

Record shapes, dtypes and dynamic/static choices beside observed tracing.
Separate Python trace diagnostics from compilation evidence.
Keep ordinary numerical hyperparameters dynamic when their values need not determine program structure.
Compare per-step host consumption with one final host consumption while checking equivalent results.

The problem

A new batch size makes one call pause; reading a metric makes another loop wait. Those delays can look similar, but they have different causes. We’ll record the shapes, dtypes and static choices of each call, then inspect where Python asks for device values. The evidence will guide the change.

The idea

Slow steps can come from preparing new program specializations or from waiting for already-compiled work. Use a controlled call ledger and timing boundaries to separate those causes.

Distinguish new tracing from host waiting

Change only array values while keeping shapes, dtypes and static choices fixed. Then change one signature component and record the observed trace behavior. Changing everything together makes the cause ambiguous.

Inspect host reads and logging next. Converting device results into Python values can force a wait even when no retracing occurs. That delay needs a different intervention from reducing signature variation.

Trace counts describe observed Python tracing. They are not a universal count of compiler work or a guarantee of cache reuse after restart. Combine the ledger with a trace of completed work before deciding which change improves the actual workload.

\text{Arithmetic Intensity (AI)} = \frac{W_{\text{FLOPs}}}{Q_{\text{bytes}}}, \qquad \text{Perf} \le \min(\pi_{\text{peak}},\; \beta_{\text{BW}} \cdot \text{AI})

A call signature is different from its numerical data

Our score is the sum of squared scaled inputs. With $x=[0, \ldots, 7]$ and gain $1$ , the answer is $140$ ; with gain $2$ it becomes $560$ . Passing a float32 scalar gain keeps its abstract shape and dtype unchanged when the value changes. The transformed computation can read that value at execution time.

An array of twelve elements has a different shape from the eight-element array. Ordinary jit specializes to these shapes, so this is a plausible cause of new staging work. A ledger records the actual calls before attributing a pause. Sharding, compiler options, callable identity and other context can matter too; the ledger is a diagnosis aid, not a full cache-key specification.

call → (shape, dtype, static choices) → specialization
             dynamic values → execution of that program

Observe tracing without mislabeling compilation

The score function appends its input signature to a Python list. This is a deliberate authoring diagnostic: Python runs while JAX traces the function, whereas the compiled numerical operations can run later without replaying that Python append. Never use this list as optimizer state or a reliable callback from accelerator execution.

We log the change in list length around each call, but assert numerical answers rather than a fixed event count. Cache state and implementation details affect observations. A trace event does not prove that a new executable was compiled, and absence of an event is not proof that no compiler activity occurred. Enable public jax_log_compiles diagnostics or obtain a profiler trace for stronger evidence, recording exactly which log event supports a claim.

Capture public compilation diagnostics on macOS/Linux

POSIX shell

# Capture public compilation diagnostics on macOS/Linux
JAX_LOG_COMPILES=1 python diagnosis.py 2> compile.log

Expected:

The usual ledger appears in the terminal. compile.log contains JAX diagnostics, including tracing and compilation messages when observed. Read the named stage in each line; do not count every line as a compile.

Capture diagnostics in PowerShell

PowerShell

# Capture diagnostics in PowerShell
$env:JAX_LOG_COMPILES="1"
python diagnosis.py 2> compile.log
Remove-Item Env:JAX_LOG_COMPILES

Expected:

The ledger appears in the terminal and diagnostics are captured in compile.log. Removing the environment variable restores the shell setting after the run.

Choose static arguments for structure rather than convenience

A Boolean selecting one of two Python branches can be static so tracing selects the branch. Its value then becomes part of specialization, so changing that choice can trigger staging. Static arguments must be hashable and are unsuitable for arbitrary numerical array data. Avoid declaring a rapidly changing scalar gain static just to make Python logic convenient.

If the gain only participates in arithmetic, pass it as a typed scalar array as we did. If a runtime Boolean chooses arithmetic, express the selection with lax.cond or another appropriate array operation. Static metadata and dynamic numerical values serve different purposes; changing one into the other can change the number of specializations and the program itself.

A host metric is an explicit synchronization boundary

float(score) needs the scalar contents in Python. That read waits for any outstanding work needed to produce it. Keeping a scalar loss as a JAX array allows later device computations to consume it without Python inspecting its value. A chain can then be synchronized once at the final output.

The exercise compares six dependent updates with host consumption after each step against six updates with a final wait. Both variants must produce the same scalar. Their wall-clock boundaries include different waits and conversion overhead; we record both but assert no speed ordering. Removing intermediate logging also removes observability, so production systems may use bounded logging intervals instead of silently dropping metrics.

A new value is not necessarily a new trace

Predict before looking: Which call changes the input signature?

Read the figure

Each bar reports the additional Python tracing events observed during one call category. The sequence is $(1,0,0,1)$ : the initial call traces, new array values do not, a new dynamic gain does not, and a new input shape triggers another trace.

The two zero-height bars are meaningful observations. The calls still execute and can return changed values; zero means the trace counter did not increase for those calls.

Connect it to the computation

The existing specialization can accept compatible dynamic values without rerunning the traced Python body. Changing the example’s shape from length $8$ to length $12$ changes the abstract input signature, so a new trace is observed. This isolates shape from ordinary value changes.

The vertical axis counts Python traces, not seconds, device kernels, or an exact count of backend compilations. To diagnose runtime overhead, combine this kind of counter with synchronized measurements or a profiler. A changed answer alone is not evidence of retracing.

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/08-performance/02-diagnose-recompilation-and-host-synchronization/code/main.py (or downloaded exercises/performance-02.py / exercises/performance-02.ipynb) and execute from the repository root with python3 phases/08-performance/02-diagnose-recompilation-and-host-synchronization/code/main.py.

1 Prepare a trace ledger

Create diagnosis.py in your activated course environment and add this block.

main.py · append build step 1

# Step 1 — Prepare a trace ledger: This CPU experiment deliberately keeps the trace ledger outside...
import json
import time
import numpy as np
import jax
import jax.numpy as jnp
# Query the active JAX devices into `cpu`.
cpu = jax.devices("cpu")[0]
trace_events = []
def score(x, gain):
    # Authoring diagnostic: this Python effect happens during tracing.
    # It is NOT model state and does not count executable compilations.
    trace_events.append({"shape": list(x.shape), "dtype": str(x.dtype)})
    return jnp.sum((x * gain)**2)
# Wrap with `jax.jit` (`compiled_score`) so XLA traces and compiles the function.
compiled_score = jax.jit(score)
x8 = jax.device_put(np.arange(8, dtype=np.float32), cpu)
x12 = jax.device_put(np.arange(12, dtype=np.float32), cpu)
# Synchronize host execution until asynchronous device computation completes.
jax.block_until_ready((x8, x12))

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.
Mesh + PartitionSpec + NamedSharding — Maps logical tensor axes onto physical device mesh axes for SPMD data, tensor, or pipeline parallelism.
jax.block_until_ready(output) — Synchronizes with the accelerator/CPU device so asynchronous dispatch finishes before wall-clock timing.

This CPU experiment deliberately keeps the trace ledger outside numerical state. Run it in a fresh process when comparing observations.

2 Record call signatures and observed traces

Append this block to the same file. Run the assembled file with python diagnosis.py.

main.py · append build step 2

# Step 2 — Record call signatures and observed traces: The ledger relates each controlled signature change to the...
call_rows = []
def observed_call(label, x, gain):
    before = len(trace_events)
    # Place `result` explicitly onto the target JAX device.
    result = compiled_score(x, jax.device_put(np.float32(gain), cpu))
    # Synchronize host execution until asynchronous device computation completes.
    result.block_until_ready()
    # Convert `expected` to a host NumPy array for inspection or verification.
    expected = np.sum((np.asarray(x) * np.float32(gain))**2)
    # Convert `` to a host NumPy array for inspection or verification.
    np.testing.assert_allclose(np.asarray(result), expected, rtol=2e-5, atol=2e-5)
    call_rows.append({"label": label, "shape": list(x.shape), "dtype": str(x.dtype),
                      "observed_trace_delta": len(trace_events) - before,
                      "result": float(result)})
    return result

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).
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
Mesh + PartitionSpec + NamedSharding — Maps logical tensor axes onto physical device mesh axes for SPMD data, tensor, or pipeline parallelism.
jax.block_until_ready(output) — Synchronizes with the accelerator/CPU device so asynchronous dispatch finishes before wall-clock timing.

The ledger relates each controlled signature change to the observed Python event delta, with an independent NumPy value check.

3 Run the controlled call sequence

Append this block to the same file. Run the assembled file with python diagnosis.py.

main.py · append build step 3

# Step 3 — Run the controlled call sequence: Inspect the event deltas on your own environment.
observed_call("initial signature", x8, 1.)
observed_call("same signature changed data", x8 + 1., 1.)
observed_call("same signature changed dynamic gain", x8, 2.)
observed_call("different input shape", x12, 1.)
print(json.dumps({"jax": jax.__version__, "backend": str(cpu),
                  "calls": call_rows, "python_trace_events": trace_events}, indent=2))
assert call_rows[0]["result"] == 140.
assert call_rows[2]["result"] == 560.
assert call_rows[3]["result"] == 506.

Code building blocks in this step

observed_call(...) — Call observed_call with your updated parameters or inputs from this lesson's workspace.
json.dumps(...) — Call json.dumps with your updated parameters or inputs from this lesson's workspace.
assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

Inspect the event deltas on your own environment. The required outputs are checked numeric scores, not a fabricated trace or compilation count.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build

# Step 1 — Prepare a trace ledger: This CPU experiment deliberately keeps the trace ledger outside...
import json
import time
import numpy as np
import jax
import jax.numpy as jnp
# Query the active JAX devices into `cpu`.
cpu = jax.devices("cpu")[0]
trace_events = []
def score(x, gain):
    # Authoring diagnostic: this Python effect happens during tracing.
    # It is NOT model state and does not count executable compilations.
    trace_events.append({"shape": list(x.shape), "dtype": str(x.dtype)})
    return jnp.sum((x * gain)**2)
# Wrap with `jax.jit` (`compiled_score`) so XLA traces and compiles the function.
compiled_score = jax.jit(score)
x8 = jax.device_put(np.arange(8, dtype=np.float32), cpu)
x12 = jax.device_put(np.arange(12, dtype=np.float32), cpu)
# Synchronize host execution until asynchronous device computation completes.
jax.block_until_ready((x8, x12))

# Step 2 — Record call signatures and observed traces: The ledger relates each controlled signature change to the...
call_rows = []
def observed_call(label, x, gain):
    before = len(trace_events)
    # Place `result` explicitly onto the target JAX device.
    result = compiled_score(x, jax.device_put(np.float32(gain), cpu))
    # Synchronize host execution until asynchronous device computation completes.
    result.block_until_ready()
    # Convert `expected` to a host NumPy array for inspection or verification.
    expected = np.sum((np.asarray(x) * np.float32(gain))**2)
    # Convert `` to a host NumPy array for inspection or verification.
    np.testing.assert_allclose(np.asarray(result), expected, rtol=2e-5, atol=2e-5)
    call_rows.append({"label": label, "shape": list(x.shape), "dtype": str(x.dtype),
                      "observed_trace_delta": len(trace_events) - before,
                      "result": float(result)})
    return result

# Step 3 — Run the controlled call sequence: Inspect the event deltas on your own environment.
observed_call("initial signature", x8, 1.)
observed_call("same signature changed data", x8 + 1., 1.)
observed_call("same signature changed dynamic gain", x8, 2.)
observed_call("different input shape", x12, 1.)
print(json.dumps({"jax": jax.__version__, "backend": str(cpu),
                  "calls": call_rows, "python_trace_events": trace_events}, indent=2))
assert call_rows[0]["result"] == 140.
assert call_rows[2]["result"] == 560.
assert call_rows[3]["result"] == 506.

Run the example

python3 phases/08-performance/02-diagnose-recompilation-and-host-synchronization/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

# Step 1 — Prepare a trace ledger: This CPU experiment deliberately keeps the trace ledger outside...
import json
import time
import numpy as np
import jax
import jax.numpy as jnp
# Query the active JAX devices into `cpu`.
cpu = jax.devices("cpu")[0]
trace_events = []
def score(x, gain):
    # Authoring diagnostic: this Python effect happens during tracing.
    # It is NOT model state and does not count executable compilations.
    trace_events.append({"shape": list(x.shape), "dtype": str(x.dtype)})
    return jnp.sum((x * gain)**2)
# Wrap with `jax.jit` (`compiled_score`) so XLA traces and compiles the function.
compiled_score = jax.jit(score)
x8 = jax.device_put(np.arange(8, dtype=np.float32), cpu)
x12 = jax.device_put(np.arange(12, dtype=np.float32), cpu)
# Synchronize host execution until asynchronous device computation completes.
jax.block_until_ready((x8, x12))
# Step 2 — Record call signatures and observed traces: The ledger relates each controlled signature change to the...
call_rows = []
def observed_call(label, x, gain):
    before = len(trace_events)
    # Place `result` explicitly onto the target JAX device.
    result = compiled_score(x, jax.device_put(np.float32(gain), cpu))
    # Synchronize host execution until asynchronous device computation completes.
    result.block_until_ready()
    # Convert `expected` to a host NumPy array for inspection or verification.
    expected = np.sum((np.asarray(x) * np.float32(gain))**2)
    # Convert `` to a host NumPy array for inspection or verification.
    np.testing.assert_allclose(np.asarray(result), expected, rtol=2e-5, atol=2e-5)
    call_rows.append({"label": label, "shape": list(x.shape), "dtype": str(x.dtype),
                      "observed_trace_delta": len(trace_events) - before,
                      "result": float(result)})
    return result
# Step 3 — Run the controlled call sequence: Inspect the event deltas on your own environment.
observed_call("initial signature", x8, 1.)
observed_call("same signature changed data", x8 + 1., 1.)
observed_call("same signature changed dynamic gain", x8, 2.)
observed_call("different input shape", x12, 1.)
print(json.dumps({"jax": jax.__version__, "backend": str(cpu),
                  "calls": call_rows, "python_trace_events": trace_events}, indent=2))
assert call_rows[0]["result"] == 140.
assert call_rows[2]["result"] == 560.
assert call_rows[3]["result"] == 506.

Expected result: A JSON ledger of four calls and observed Python trace events. Checked scores include $140$ , $560$ and $506$ ; event deltas are recorded rather than fixed by the lesson.

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:47:35.719968+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

{
  "jax": "0.9.2",
  "backend": "TFRT_CPU_0",
  "calls": [
    {
      "label": "initial signature",
      "shape": [
        8
      ],
      "dtype": "float32",
      "observed_trace_delta": 1,
      "result": 140.0
    },
    {
      "label": "same signature changed data",
      "shape": [
        8
      ],
      "dtype": "float32",
      "observed_trace_delta": 0,
      "result": 204.0
    },
    {
      "label": "same signature changed dynamic gain",
      "shape": [
        8
      ],
      "dtype": "float32",
      "observed_trace_delta": 0,
      "result": 560.0
    },
    {
      "label": "different input shape",
      "shape": [
        12
      ],
      "dtype": "float32",
      "observed_trace_delta": 1,
      "result": 506.0
    }
  ],
  "python_trace_events": [
    {
      "shape": [
        8
      ],
      "dtype": "float32"
    },
    {
      "shape": [
        12
      ],
      "dtype": "float32"
    }
  ]
}
{
  "jax": "0.9.2",
  "backend": "TFRT_CPU_0",
  "calls": [
    {
      "label": "initial signature",
      "shape": [
        8
      ],
      "dtype": "float32",
      "observed_trace_delta": 1,
      "result": 140.0
    },
    {
      "label": "same signature changed data",
      "shape": [
        8
      ],
      "dtype": "float32",
      "observed_trace_delta": 0,
      "result": 204.0
    },
    {
      "label": "same signature changed dynamic gain",
      "shape": [
        8
      ],
      "dtype": "float32",
      "observed_trace_delta": 0,
      "result": 560.0
    },
    {
      "label": "different input shape",
      "shape": [
        12
      ],
      "dtype": "float32",
      "observed_trace_delta": 1,
      "result": 506.0
    }
  ],
  "python_trace_events": [
    {
      "shape": [
        8
      ],
      "dtype": "float32"
    },
    {
      "shape": [
        12
      ],
      "dtype": "float32"
    }
  ]
}
Static branch trace observations: [{'shape': [8], 'square': True}, {'shape': [8], 'square': False}]
Per-step host reads seconds: 0.00012137508019804955
Final-only wait seconds: 3.5999808460474014e-05
Added ledger row: {'label': 'new length sixteen', 'shape': [16], 'dtype': 'float32', 'observed_trace_delta': 1, 'result': 1240.0}
Masked versus unmasked padded objective: 55.0 58.0
Observed runtime-Python-branch tracing failure
PASS: performance-02

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

Contrast static branch metadata with a dynamic gain

Contrast static branch metadata with a dynamic gain · Python experiment

# Experiment — Contrast static branch metadata with a dynamic gain: The static Boolean controls Python program structure.
branch_events = []
def branch_score(x, square):
    branch_events.append({"shape": list(x.shape), "square": square})
    return jnp.sum(x**2) if square else jnp.sum(x)
# Wrap with `jax.jit` (`static_score`) so XLA traces and compiles the function.
static_score = jax.jit(branch_score, static_argnames=("square",))
for square in [True, True, False]:
    out = static_score(x8, square=square)
    # Convert `` to a host NumPy array for inspection or verification.
    np.testing.assert_allclose(np.asarray(out), 140. if square else 28.)
print("Static branch trace observations:", branch_events)

Expected: Correct scores for both branches, plus the trace events observed in this process.

The static Boolean controls Python program structure. The earlier numerical gain was dynamic arithmetic data. These are distinct specialization choices; the printed diagnostic is not an executable compilation count.

Keep a metric on device until the reporting boundary

Keep a metric on device until the reporting boundary · Python experiment

# Experiment — Keep a metric on device until the reporting boundary: The geometric-series result checks the computation independently.
# Wrap with `jax.jit` (`update`) so XLA traces and compiles the function.
update = jax.jit(lambda value: 0.9 * value + 1.)
# Place `initial` explicitly onto the target JAX device.
initial = jax.device_put(np.float32(0.), cpu)
# Synchronize host execution until asynchronous device computation completes.
update(initial).block_until_ready()
def metric_run(read_each):
    value = initial
    observed = []
    start = time.perf_counter()
    for _ in range(6):
        value = update(value)
        if read_each:
            observed.append(float(value))
    # Synchronize host execution until asynchronous device computation completes.
    value.block_until_ready()
    elapsed = time.perf_counter() - start
    return value, elapsed, observed
read_result, read_seconds, metrics = metric_run(True)
final_result, final_seconds, _ = metric_run(False)
reference_value = 10. * (1. - 0.9**6)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(read_result), reference_value, rtol=2e-5)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(final_result), reference_value, rtol=2e-5)
assert len(metrics) == 6
print("Per-step host reads seconds:", read_seconds)
print("Final-only wait seconds:", final_seconds)

Expected: Equivalent final values near $4.68559$ and two measured intervals; no required speed ordering.

The geometric-series result checks the computation independently. Moving host consumption changes the measurement/observability boundary, so keep that difference beside both times.

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/08-performance/02-diagnose-recompilation-and-host-synchronization/code/main.py (or inside the Make it yours cell of exercises/performance-02.ipynb) and run python3 phases/08-performance/02-diagnose-recompilation-and-host-synchronization/code/main.py to verify your prediction before opening the reference solution.

Add a call with sixteen float32 elements. Calculate its expected score independently and extend the signature ledger without labeling a trace event as a compilation.

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:

Mesh + PartitionSpec + NamedSharding — Maps logical tensor axes onto physical device mesh axes for SPMD data, tensor, or pipeline parallelism.

2. Step-by-step implementation plan:

Set up the changed input or configuration parameters requested in the prompt.
Run the target function or transformation and assert that the output matches your prediction.

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

# Exercise solution: Add a call with sixteen float32 elements.
x16 = jax.device_put(...)  # TODO: compute x16
observed_call("new length sixteen", x16, 1.)
assert call_rows[-1]["result"]  # TODO: complete assertion check
print("Added ledger row:", call_rows[-1])

Reveal the reference solution

Reference solution

# Exercise solution: Add a call with sixteen float32 elements.
x16 = jax.device_put(np.arange(16, dtype=np.float32), cpu)
observed_call("new length sixteen", x16, 1.)
assert call_rows[-1]["result"] == sum(i*i for i in range(16))
print("Added ledger row:", call_rows[-1])

Pad a final batch while preserving the objective

Transfer

A final batch has five values while other batches have eight. Pad to eight and use an explicit validity mask in a dynamic array. Compare the masked sum of squares to an ordinary five-value calculation. Explain why padding without masking changes objectives with offsets.

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:

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.
Mesh + PartitionSpec + NamedSharding — Maps logical tensor axes onto physical device mesh axes for SPMD data, tensor, or pipeline parallelism.

2. Step-by-step implementation plan:

Wrap with jax.jit (masked_score) so XLA traces and compiles the function.
Place padded explicitly onto the target JAX device.
Combine or mask array elements to form masked.
Convert to a host NumPy array for inspection or verification.
Aggregate array values to compute wrong.

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

# Pad a final batch while preserving the objective (Transfer): The mask removes three synthetic rows.
# Wrap with `jax.jit` (`masked_score`) so XLA traces and compiles the function.
masked_score = jax.jit(...)  # TODO: compute masked_score
five = np.arange(...)  # TODO: compute five
# Place `padded` explicitly onto the target JAX device.
padded = jax.device_put(...)  # TODO: compute padded
mask = jax.device_put(...)  # TODO: compute mask
# Combine or mask array elements to form `masked`.
masked = masked_score(...)  # TODO: compute masked
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(masked), sum((i+1)**2 for i in range(5)))
# Aggregate array values to compute `wrong`.
wrong = jnp.sum(...)  # TODO: compute wrong
assert float(wrong)  # TODO: complete assertion check
print("Masked versus unmasked padded objective:", float(masked), float(wrong))

Show a hint

Use a fixed-size input and mask; compute the offset transformation before masking.

Reveal reference solution and reasoning

Pad a final batch while preserving the objective · Python solution

# Pad a final batch while preserving the objective (Transfer): The mask removes three synthetic rows.
# Wrap with `jax.jit` (`masked_score`) so XLA traces and compiles the function.
masked_score = jax.jit(lambda values, mask: jnp.sum(jnp.where(mask, (values + 1.)**2, 0.)))
five = np.arange(5, dtype=np.float32)
# Place `padded` explicitly onto the target JAX device.
padded = jax.device_put(np.pad(five, (0, 3)), cpu)
mask = jax.device_put(np.arange(8) < 5, cpu)
# Combine or mask array elements to form `masked`.
masked = masked_score(padded, mask)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(masked), sum((i+1)**2 for i in range(5)))
# Aggregate array values to compute `wrong`.
wrong = jnp.sum((padded + 1.)**2)
assert float(wrong) == float(masked) + 3.
print("Masked versus unmasked padded objective:", float(masked), float(wrong))

The mask removes three synthetic rows. Zero padding alone adds three ones after the offset, so a fixed shape is useful only if it retains the original mathematical objective.

Reproduce and repair runtime Python branching

Diagnosis

Write a jitted function using if gain > $1$ on a scalar array. Capture its tracing failure and replace the branch with lax.cond. Verify both branches against independent arithmetic.

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:

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.
Mesh + PartitionSpec + NamedSharding — Maps logical tensor axes onto physical device mesh axes for SPMD data, tensor, or pipeline parallelism.

2. Step-by-step implementation plan:

Reproduce and repair runtime Python branching (Diagnosis): The failure identifies a mismatch between Python control...
Run the boundary check and catch the expected exception:
Wrap with jax.jit (repaired) so XLA traces and compiles the function.
Place result explicitly onto the target JAX device.
Convert to a host NumPy array for inspection or verification.

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

# Reproduce and repair runtime Python branching (Diagnosis): The failure identifies a mismatch between Python control...
def broken_branch(values, gain):
    if gain > 1.:
        return ...  # TODO: return computed result
    return ...  # TODO: return computed result
# Run the boundary check and catch the expected exception:
try:
    jax.jit(broken_branch)(x8, jax.device_put(np.float32(2.), cpu))
except jax.errors.TracerBoolConversionError:
    print("Observed runtime-Python-branch tracing failure")
else:
    raise AssertionError("runtime branch unexpectedly accepted")
# Wrap with `jax.jit` (`repaired`) so XLA traces and compiles the function.
repaired = jax.jit(...)  # TODO: compute repaired
for gain, expected in [(0.5, 28.), (2., 140.)]:
    # Place `result` explicitly onto the target JAX device.
    result = repaired(...)  # TODO: compute result
    # Convert `` to a host NumPy array for inspection or verification.
    np.testing.assert_allclose(np.asarray(result), expected)

Show a hint

A runtime scalar predicate is a tracer, so Python cannot decide its truth value during tracing.

Reveal reference solution and reasoning

Reproduce and repair runtime Python branching · Python solution

# Reproduce and repair runtime Python branching (Diagnosis): The failure identifies a mismatch between Python control...
def broken_branch(values, gain):
    if gain > 1.:
        return jnp.sum(values**2)
    return jnp.sum(values)
# Run the boundary check and catch the expected exception:
try:
    jax.jit(broken_branch)(x8, jax.device_put(np.float32(2.), cpu))
except jax.errors.TracerBoolConversionError:
    print("Observed runtime-Python-branch tracing failure")
else:
    raise AssertionError("runtime branch unexpectedly accepted")
# Wrap with `jax.jit` (`repaired`) so XLA traces and compiles the function.
repaired = jax.jit(lambda values, gain: jax.lax.cond(gain > 1., lambda v:jnp.sum(v**2), lambda v:jnp.sum(v), values))
for gain, expected in [(0.5, 28.), (2., 140.)]:
    # Place `result` explicitly onto the target JAX device.
    result = repaired(x8, jax.device_put(np.float32(gain), cpu))
    # Convert `` to a host NumPy array for inspection or verification.
    np.testing.assert_allclose(np.asarray(result), expected)

The failure identifies a mismatch between Python control flow and runtime numerical data. lax.cond preserves a dynamic choice; declaring gain static would choose a different specialization contract instead.

Diagnose the result

If pauses correlate with varying batch lengths, record their signatures and check an explicitly masked fixed-size alternative. If host reads appear inside every update, first specify the reporting boundary and verify numerical equivalence before comparing synchronized observations.

Carry forward

Record observed staging diagnostics without inventing compilation counts.
Static metadata changes program specialization; ordinary arithmetic values can remain dynamic.
Padding needs an objective-preserving mask.
Host metrics provide useful observability and introduce completion boundaries.

Keep your evidence

Download evidence template

Keep the controlled call ledger, actual trace observations labeled as traces, static/dynamic explanation, observed host-consumption intervals, masked-versus-unmasked counterexample and reproduced/repaired tracing failure.

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).

Diagnose recompilation and host synchronization

What you will be able to do

The problem

The idea

Distinguish new tracing from host waiting

A call signature is different from its numerical data

Observe tracing without mislabeling compilation

Capture public compilation diagnostics on macOS/Linux

Capture diagnostics in PowerShell

Choose static arguments for structure rather than convenience

A host metric is an explicit synchronization boundary

A new value is not necessarily a new trace

Read the figure

Connect it to the computation

1 Prepare a trace ledger

2 Record call signatures and observed traces

3 Run the controlled call sequence

Run the example

Recorded reference execution

Contrast static branch metadata with a dynamic gain

Keep a metric on device until the reporting boundary

Make it yours

Pad a final batch while preserving the objective

Reproduce and repair runtime Python branching

Check your understanding

Diagnose the result

Carry forward

Keep your evidence

Primary references