Tracing, static arguments and recompiles

Almost every confusing jit error is a tracing error; this mental model explains and fixes them.

Under jit, arguments become tracers: stand-ins that know shape and dtype (ShapedArray(float32[3])) but not values. Your Python code runs once on them. Anything that needs a concrete value fails:

if x > 0: or while x < 10: on a traced array raises TracerBoolConversionError.
jnp.zeros(n) with traced n fails because output shapes must be known at compile time.
int(x) or x.item() raises a ConcretizationTypeError.

Fixes, in order of preference:

Use array-level logic: jnp.where, jnp.maximum, masks.
Use structured control flow from jax.lax (next lesson).
Mark the argument static: jax.jit(f, static_argnames="n"). A static argument is treated as a compile-time constant and must be hashable; each new value compiles a new version.

Know what triggers recompilation: a new input shape or dtype, a new static argument value, or a new pytree structure. Variable-length inputs (sentences, batches at the end of an epoch) can therefore cause a compile per length; pad to a few bucket sizes instead.

Shape information is static, so x.shape[0] is a plain Python int you can branch on freely. Python loops over a fixed range(3) are also fine: they simply unroll into the traced program, which can make compile time grow for long loops.

Example

import jax
import jax.numpy as jnp

@jax.jit
def bad(x):
    if x.sum() > 0:          # Python bool of a tracer
        return x
    return -x

try:
    bad(jnp.ones(3))
except jax.errors.TracerBoolConversionError:
    print("TracerBoolConversionError")

@jax.jit
def good(x):
    return jnp.where(x.sum() > 0, x, -x)
print(good(jnp.ones(3)))                 # [1. 1. 1.]

rep = jax.jit(lambda x, n: jnp.tile(x, n), static_argnames="n")  # n sets the shape
print(rep(jnp.array([1, 2]), n=3))       # [1 2 1 2 1 2]

print(jax.jit(lambda x: x.shape[0] * 2)(jnp.ones((5, 2))))  # 10: shapes are static

Remember

Tracers know shape and dtype, never values, so Python branching on them fails
Static arguments are compile-time constants; each new value recompiles
New shapes, dtypes or pytree structures also trigger recompilation

Watch out for

Marking a frequently changing argument (like a step counter) static, causing a recompile every step.

Check yourself

Which of these is safe inside a jitted function whose argument x is a traced array?

Official docs: docs.jax.dev/en/latest/201/jit.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks