Skip to content

← JAX Foundations step 4 of 25

Medium End-to-End

Tracer Errors: Boolean Conversion and Leaks

Jump to the task

Two tracer problems with different fixes

A tracer represents a value while JAX transforms a function. In an ordinary JIT trace, a dynamic array's shape and dtype are known, but its numerical values are not available for Python to inspect.

Asking Python to convert a traced condition to bool is a conversion error. Saving a tracer somewhere outside the transformed function is a leak. These are different problems, not old and new names for the same error.

Replace a value-dependent Python branch

import jax
import jax.numpy as jnp

def clip_first(x):
    if x[0] < 0:
        return x.at[0].set(0)
    return x

# jax.jit(clip_first)(jnp.array([-2., 3.]))
# raises TracerBoolConversionError

Python needs to choose the branch while tracing, but x[0] is dynamic. JAX can represent runtime decisions; use a JAX operation to express them. For clipping every negative entry, an elementwise selection fits:

def clip_all(x):
    return jnp.where(x < 0, 0, x)

result = jax.jit(clip_all)(jnp.array([3., -7., 0.5]))
# [3., 0., 0.5]

The mask is [False, True, False]. where selects zero where the mask is true and the original value elsewhere. It keeps the array shape. Unlike clip_first, this checks every entry. For a whole-function decision based on one scalar condition, consider jax.lax.cond instead; cond is not an elementwise mask operation.

A loop over range(len(x)) is not itself a leak: the leading dimension is static under ordinary JIT. A Python if on x[i] inside that loop still has the boolean-conversion problem. Other requests for concrete values, such as tolist(), can raise different conversion errors.

A real leak saves a tracer outside its scope

saved = []

@jax.jit
def leaky(x):
    y = x + 1
    saved.append(y)  # writes a tracer into an outer list during tracing
    return y

leaky(jnp.array([1., 2.]))
# saved[0] + 1 would raise UnexpectedTracerError

Return needed values normally, then store the returned concrete array in host code. Merely renaming the list or replacing if with where does not fix an escaped tracer.

with jax.checking_leaks(): around a fresh transformed call can report a leak when the trace finishes, closer to its source than later use. This debugging check has overhead and can report false positives with Python debuggers. Local temporary containers are not automatically leaks; the problem is letting traced values escape through side effects.

Problem

Implement safe_clip_positive(x) to replace each negative entry with zero, leaving all nonnegative entries unchanged. Despite the function name, zero is allowed: the result is nonnegative, not strictly positive.

X is a one-dimensional array of finite float32 values. Return a JAX array with the same shape and dtype, without modifying x. An empty input stays empty. [3, -7, 0.5] becomes [3, 0, 0.5].

The function must work under jax.jit. Use an elementwise JAX operation such as jnp.where or jnp.maximum, without a Python if on array values or saving intermediate tracers to outside state.