# Tracing, static arguments, and recompilation
Phase 02: JAX transformations · about 60 minutes · CPU
## 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)
$$
### Pause and reason
A loop passes a new Python configuration value on every call as a static argument. What should you inspect?
Compare your reasoning
Inspect whether that value truly changes program structure. If it is ordinary changing data, represent it dynamically where possible. Keep the call ledger and compare traces after changing one factor at a time.
Reference solution
```python
# 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")
```
Hint
`jnp.where` selects elementwise; both candidate arrays are computed expressions. It is not a Python branch per element.
Reference solution and reasoning
```python
# 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.
Hint
Do not use elapsed time alone to infer whether tracing occurred. Values and shape are different inputs to your reasoning.
Reference solution and reasoning
```python
# 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.
Answer and explanation
Express the runtime decision with compatible JAX control flow.
A runtime decision belongs in the staged computation. Static configuration is useful for a small program choice, but making changing data static can create unnecessary specializations.