Common jit Errors and Fixes

Let's walk through the most common errors you'll encounter with JIT and how to fix each one.

Error 1: ConcretizationTypeError

import jax
import jax.numpy as jnp

# Problem: Python if with traced value
@jax.jit
def categorize(x):
    if x > 0:       # ERROR!
        return 1
    elif x < 0:
        return -1
    else:
        return 0

# Fix: use jnp.sign (or jnp.where for custom logic)
@jax.jit
def categorize_fixed(x):
    return jnp.sign(x).astype(jnp.int32)

Error 2: Shapes Must Be Static

import jax
import jax.numpy as jnp

# Problem: dynamic output shape
@jax.jit
def filter_positive(x):
    return x[x > 0]  # Output size depends on values — not allowed!

# Fix 1: Use a fixed-size output with padding
@jax.jit
def filter_positive_fixed(x, max_size):
    mask = x > 0
    indices = jnp.where(mask, size=max_size, fill_value=0)
    return x[indices[0]]

# Fix 2: Don't JIT this particular function
# (sometimes the simplest answer)
def filter_positive_eager(x):
    return x[x > 0]  # Works fine without jit

Error 3: Python for-Loop Over Traced Values

import jax
import jax.numpy as jnp

# Problem: Python range with traced value
@jax.jit
def repeat_add(x, n):
    result = x
    for i in range(n):  # ERROR if n is traced
        result = result + x
    return result

# Fix 1: make n static
from functools import partial

@partial(jax.jit, static_argnums=(1,))
def repeat_add_static(x, n):
    result = x
    for i in range(n):  # n is concrete now — loop is unrolled
        result = result + x
    return result

# Fix 2: use jax.lax.fori_loop
@jax.jit
def repeat_add_functional(x, n):
    return jax.lax.fori_loop(0, n, lambda i, r: r + x, x)

When NOT to JIT

• Debugging: Run without JIT first to get normal Python errors and print output.
• One-off computations: If you call a function once, the compilation cost outweighs the benefit.
• Highly dynamic shapes: If input shapes change every call, JIT will recompile every time, making it slower than eager mode.
• Very small computations: The overhead of dispatching to XLA can exceed the computation time for tiny arrays.