# Diagnose recompilation and host synchronization
Phase 08: Performance diagnosis · about 80 minutes · CPU
## 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})
$$
### Pause and reason
If the trace count stays constant but a step is slow, what should you inspect?
Compare your reasoning
Inspect synchronization, transfers, data preparation and logging. A stable trace count does not rule out those boundaries.
Reference solution
```python
# 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])
```
Hint
Use a fixed-size input and mask; compute the offset transformation before masking.
Reference solution and reasoning
```python
# 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.
Hint
A runtime scalar predicate is a tracer, so Python cannot decide its truth value during tracing.
Reference solution and reasoning
```python
# 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.
Answer and explanation
An observed Python tracing event, which is not an executable compilation count
Python executes during tracing; compiled array computation may run repeatedly without that append. Correlate with public compiler logs or traces before claiming recompilation.