Debugging Impure Code: Error Messages Decoded

When you violate purity inside a JIT-compiled (or grad/vmap) function, JAX throws specific errors. Learning to read these error messages will dramatically speed up your debugging.

ConcretizationTypeError

This is the most common JAX error. It means you're trying to use a traced value (an abstract placeholder) in a context that requires a concrete Python value.

import jax
import jax.numpy as jnp

@jax.jit
def bad_fn(x):
    if x > 0:  # Needs concrete bool, but x is a tracer
        return x
    else:
        return -x

# ConcretizationTypeError: Abstract tracer value encountered
# where a concrete value is expected: Traced<ShapedArray(float32[])>
# This often arises when using Python control flow with traced values.

# Fix 1: Use jnp.where (no branching)
@jax.jit
def fix1(x):
    return jnp.where(x > 0, x, -x)

# Fix 2: Use jax.lax.cond (traced branching)
@jax.jit
def fix2(x):
    return jax.lax.cond(x > 0, lambda x: x, lambda x: -x, x)

TracerArrayConversionError

This error occurs when you try to convert a JAX tracer to a NumPy array or Python scalar during tracing.

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

@jax.jit
def bad_conversion(x):
    # Trying to use a traced value as a Python int
    n = int(x.shape[0])  # This works (shape is static)
    val = int(x[0])       # ERROR: can't convert tracer to int
    return x[:val]

# Fix: keep everything in JAX-land
@jax.jit
def fixed(x):
    return jax.lax.dynamic_slice(x, (0,), (2,))

Leaked Tracer Error

This happens when a tracer escapes the scope of a transformation — usually by being stored in a global variable or external data structure.

import jax
import jax.numpy as jnp

stored_values = []

@jax.jit
def leaky_fn(x):
    y = x + 1
    stored_values.append(y)  # Tracer escapes into global list!
    return y

# This may raise a leaked tracer error or produce silent bugs
# Fix: don't store traced values outside the function

Debugging Strategy: Remove Transformations First

import jax
import jax.numpy as jnp

def my_function(x):
    """Some complex computation."""
    intermediate = jnp.sin(x) * 2
    # ... more code ...
    return jnp.sum(intermediate)

# Step 1: Test WITHOUT jit first
x = jnp.array([1.0, 2.0, 3.0])
result = my_function(x)  # If this works, the logic is correct
print(result)

# Step 2: Use print to inspect values
print(f"intermediate shape: {jnp.sin(x).shape}")

# Step 3: Add jit and check for errors
jitted = jax.jit(my_function)
result = jitted(x)

# Step 4: If jit fails, use jax.make_jaxpr to see the traced computation
print(jax.make_jaxpr(my_function)(x))