Frozen, passed, or escaped

The run below changes a closed-over array between two calls and the jitted function does not notice. First call prints ones. Mutate bias[0] to 100. Second call prints ones again, while the same expression outside jit prints 101, so there is no doubt about what the array now holds.

Nothing failed here, and nothing recompiled, which is the part worth being precise about. The trace read bias once and the value went into the recording. The second call arrived with the same shape and dtype, matched the cached executable, and never ran the Python body at all. add_bias._cache_size() says one, because a closed-over value is not part of the key that chapter 3 describes.

So the rule to carry is short. A value that changes between calls belongs in the argument list, where every call rebinds it. A value that will not change can be closed over, and the trace will keep its own copy.

Change a closed-over array and nothing recompiles, because nothing in the key changed.

import jax
import jax.numpy as jnp
import numpy as np

bias = np.zeros(3)

@jax.jit
def add_bias(x):
    return x + bias

print(add_bias(jnp.ones(3)))
bias[0] = 100.0
print(add_bias(jnp.ones(3)))
print(add_bias._cache_size())
print(jnp.ones(3) + bias)

# [1. 1. 1.]
# [1. 1. 1.]
# 1
# [101.   1.   1.]

Trace a function that closes over the Python float 2.0 and the number appears inside the equation itself: mul a 2.0. There is no separate storage for it, and no way to change it short of retracing.

Close over an array instead and the recording grows a slot in front of the arguments, with the array's value carried alongside as a constant. The .consts list on the traced object holds it, one entry, shape (2, 3) in the run below. The kernel path's source lesson at /l/source names the two slots the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → grammar uses for this; what matters here is that a copy of the value is now part of the program.

That copy has a size. A closed-over parameter tree of a gigabyte is a gigabyte of constants in the trace, held for as long as the compiled function is alive, and duplicated per signature you trace. Passing the same tree as an argument costs nothing extra, which is most of why idiomatic JAX threads parameters through the call rather than reaching for them from an enclosing scope.

import jax
import jax.numpy as jnp

scale = 2.0
weights = jnp.ones((2, 3))

print(jax.make_jaxpr(lambda x: x * scale)(jnp.ones(3)))
closed = jax.make_jaxpr(lambda x: weights @ x)(jnp.ones(3))
print(closed)
print(len(closed.consts), closed.consts[0].shape)

# { lambda ; a:f32[3]. let b:f32[3] = mul a 2.0 in (b,) }
# { lambda a:f32[2,3]; b:f32[3]. let
#     c:f32[2] = dot_general[
#       dimension_numbers=(([1], [0]), ([], []))
#       preferred_element_type=float32
#     ] a b
#   in (c,) }
# 1 (2, 3)

Append the argument to a module-level list inside a jitted function and nothing complains during the call. The append happens at trace time, the function returns its real result, and the list now holds a DynamicJaxprTracer whose trace finished a moment ago.

Touching that object is what raises, and the message is unusually generous. It names the type that escaped, states the rule it broke, names the function being traced when the leak happened, and prints the stack frames from where the value was created. Read the creation line carefully: it points at the call site, <string>:11 here, not at the append that did the storing, because the value was created when the trace began.

The last two lines hand you the follow-up before you go looking for it: an environment variable and a context manager that catch the leak earlier, which the next section runs.

import jax
import jax.numpy as jnp

seen = []

@jax.jit
def step(x):
    seen.append(x)
    return x * 2

step(jnp.ones(3))
print(type(seen[0]).__name__)
try:
    print(seen[0] + 1)
except jax.errors.UnexpectedTracerError as e:
    print(e)

# DynamicJaxprTracer
# Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type float32[3] wrapped in a DynamicJaxprTracer to escape the scope of the transformation.
# JAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.
# The function being traced when the value leaked was step at <string>:6 traced for jit.
# ------------------------------
# The leaked intermediate value was created on line <string>:11 (<module>).
# ------------------------------
# When the value was created, the final 5 stack frames (most recent last) excluding JAX-internal frames were:
# ------------------------------
# <string>:11 (<module>)
# ------------------------------
#
# To catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.
# See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.UnexpectedTracerError

Wrap the same call in jax.checking_leaks() and the failure moves from the moment you use the tracer to the moment the trace ends. That matters when the two are far apart, which they usually are: a leak stored during setup and read during evaluation gives you an error with nothing useful nearby.

What this message adds is the referrer chain. It reports the leaked tracer, then the list holding it at index 0, then the module global holding the list, so you get the path to the object that kept the reference rather than the line that used it. The same check is available without editing code, through the JAX_CHECK_TRACER_LEAKS environment variable.

It is a debugging mode rather than a setting to leave on. The check walks referrers for every trace, so it is slow, and the documentation says so plainly.

import jax
import jax.numpy as jnp

seen = []

@jax.jit
def step(x):
    seen.append(x)
    return x * 2

try:
    with jax.checking_leaks():
        step(jnp.ones(3))
except Exception as e:
    print(type(e).__name__)
    print(e)

# Exception
# Leaked trace DynamicJaxprTrace. Leaked tracer(s):
#
# Traced<ShapedArray(float32[3])>with<DynamicJaxprTrace>
# The error occurred while tracing the function step at <string>:6 for jit. This concrete value was not available in Python because it depends on the value of the argument x.
# <DynamicJaxprTracer 4624049616> is referred to by <list 4520251328>[0]
# <list 4520251328> is referred to by __main__.seen

None of this makes closing over a tracer illegal. Define an inner jitted lambda inside a jitted function, let it close over the outer function's argument, and it runs without complaint, because the inner trace is still inside the outer one.

The jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → shows what that closure turned into. The inner pjit equation takes two operands, c and d, where c is the outer function's argument and d is the array the lambda was called with. A tracer closed over inside its own trace becomes an ordinary operand of the inner call.

So the line is not about closures and not about tracers. It is about whether the trace that made the tracer is still running when you use it. Inside, the value is a normal intermediate; outside, it is a reference to a program that has already been recorded and handed off.

import jax
import jax.numpy as jnp

@jax.jit
def outer(x):
    inner = jax.jit(lambda y: y + x)
    return inner(jnp.ones(3))

print(outer(jnp.arange(3.0)))
print(jax.make_jaxpr(outer)(jnp.arange(3.0)))

# [1. 2. 3.]
# { lambda ; a:f32[3]. let
#     b:f32[3] = pjit[
#       name=outer
#       jaxpr={ lambda ; c:f32[3]. let
#           d:f32[3] = broadcast_in_dim[
#             broadcast_dimensions=()
#             shape=(3,)
#             sharding=None
#           ] 1.0
#           e:f32[3] = pjit[
#             name=<lambda>
#             jaxpr={ lambda ; f:f32[3] g:f32[3]. let h:f32[3] = add g f in (h,) }
#           ] c d
#         in (e,) }
#     ] a
#   in (b,) }