JAX Gotchas and Common Pitfalls

JAX has some behaviors that can be surprising if you’re coming from NumPy or PyTorch. This notebook covers common mistakes and how to avoid them.

Topics covered:

1. Pure functions and side effects

2. In-place updates don’t work

3. Random number generation

4. Control flow with jit

5. Tracing and concrete values

6. Floating point precision (32-bit default)

7. NaN debugging

8. Memory and recompilation

9. Pytree gotchas

10. Common error messages explained

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random, lax
import numpy as np

1. Pure Functions: No Side Effects!

JAX transformations (jit, grad, vmap) assume pure functions - same inputs always produce same outputs, and no side effects.

Gotcha: Global state changes are ignored

# BAD: Using global state
counter = 0

def bad_function(x):
    global counter
    counter += 1  # Side effect!
    return x + counter

print("Without JIT:")
print(f"  bad_function(1.0) = {bad_function(1.0)}, counter = {counter}")
print(f"  bad_function(1.0) = {bad_function(1.0)}, counter = {counter}")

counter = 0
bad_jit = jit(bad_function)

print("\nWith JIT:")
print(f"  bad_jit(1.0) = {bad_jit(1.0)}, counter = {counter}")
print(f"  bad_jit(1.0) = {bad_jit(1.0)}, counter = {counter}")
print("\n❌ The counter only incremented during tracing, not during execution!")

WARNING:2026-01-10 20:26:59,270:jax._src.xla_bridge:852: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

Without JIT:
  bad_function(1.0) = 2.0, counter = 1
  bad_function(1.0) = 3.0, counter = 2

With JIT:
  bad_jit(1.0) = 2.0, counter = 1
  bad_jit(1.0) = 2.0, counter = 1

❌ The counter only incremented during tracing, not during execution!

# GOOD: Pass state as argument, return updated state

def good_function(x, counter):
    counter = counter + 1
    return x + counter, counter

good_jit = jit(good_function)

counter = 0
print("Correct approach:")
result, counter = good_jit(1.0, counter)
print(f"  good_jit(1.0, 0) = {result}, counter = {counter}")
result, counter = good_jit(1.0, counter)
print(f"  good_jit(1.0, 1) = {result}, counter = {counter}")
print("\n✓ State is passed explicitly and returned!")

Correct approach:
  good_jit(1.0, 0) = 2.0, counter = 1
  good_jit(1.0, 1) = 3.0, counter = 2

✓ State is passed explicitly and returned!

2. In-Place Updates Don’t Work

JAX arrays are immutable. Operations like x[0] = 5 will fail.

# BAD: In-place modification
x = jnp.array([1, 2, 3])

try:
    x[0] = 10  # This will fail!
except TypeError as e:
    print(f"❌ Error: {e}")

❌ Error: JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html

# GOOD: Use .at[].set() for functional updates
x = jnp.array([1, 2, 3])

# This creates a NEW array with the modification
x_new = x.at[0].set(10)

print(f"Original: {x}")
print(f"Updated:  {x_new}")
print("\n✓ Original array is unchanged!")

Original: [1 2 3]
Updated:  [10  2  3]

✓ Original array is unchanged!

# Other .at[] operations
x = jnp.array([1, 2, 3, 4, 5])

print("Various .at[] operations:")
print(f"  x.at[0].set(10):  {x.at[0].set(10)}")
print(f"  x.at[0].add(10):  {x.at[0].add(10)}")
print(f"  x.at[0].mul(10):  {x.at[0].mul(10)}")
print(f"  x.at[0].max(10):  {x.at[0].max(10)}")
print(f"  x.at[1:3].set(0): {x.at[1:3].set(0)}")

Various .at[] operations:
  x.at[0].set(10):  [10  2  3  4  5]

  x.at[0].add(10):  [11  2  3  4  5]
  x.at[0].mul(10):  [10  2  3  4  5]

  x.at[0].max(10):  [10  2  3  4  5]

  x.at[1:3].set(0): [1 0 0 4 5]

3. Random Number Generation

JAX uses explicit PRNG keys instead of global state. This is required for reproducibility with JIT.

# BAD: Reusing the same key
key = random.PRNGKey(0)

print("Reusing same key (BAD):")
print(f"  random.normal(key): {random.normal(key)}")
print(f"  random.normal(key): {random.normal(key)}")
print("❌ Same result every time!")

Reusing same key (BAD):

  random.normal(key): 1.622642159461975
  random.normal(key): 1.622642159461975
❌ Same result every time!

# GOOD: Split the key
key = random.PRNGKey(0)

print("Splitting keys (GOOD):")
key, subkey = random.split(key)
print(f"  random.normal(subkey): {random.normal(subkey)}")
key, subkey = random.split(key)
print(f"  random.normal(subkey): {random.normal(subkey)}")
print("✓ Different results!")

Splitting keys (GOOD):

  random.normal(subkey): -2.442455768585205
  random.normal(subkey): -1.2574776411056519
✓ Different results!

# Pattern for functions that need randomness

def sample_and_transform(key, x):
    """Always pass key as first argument."""
    noise = random.normal(key, x.shape)
    return x + 0.1 * noise

# Usage
key = random.PRNGKey(42)
x = jnp.ones(3)

for i in range(3):
    key, subkey = random.split(key)
    result = sample_and_transform(subkey, x)
    print(f"Iteration {i}: {result}")

Iteration 0: [1.0605764 1.0799044 0.9091073]
Iteration 1: [0.978911   0.86372054 0.9954996 ]
Iteration 2: [0.9605102  0.9574035  0.89280283]

4. Control Flow with JIT

Python control flow (if/else, for, while) behaves differently inside JIT-compiled functions.

# BAD: Using Python if with traced values
@jit
def python_if(x):
    if x > 0:  # This fails! x is traced, can't convert to bool
        return x
    else:
        return -x

print("Python if with traced values:")
try:
    result = python_if(5.0)
    print(f"  python_if(5.0) = {result}")
except jax.errors.TracerBoolConversionError as e:
    print("  ❌ TracerBoolConversionError!")
    print("  Can't use Python if/else with traced values in JIT.")

Python if with traced values:
  ❌ TracerBoolConversionError!
  Can't use Python if/else with traced values in JIT.

# GOOD: Use lax.cond for value-dependent branches
@jit
def jax_cond(x):
    return lax.cond(
        x > 0,
        lambda x: x,      # True branch
        lambda x: -x,     # False branch
        x
    )

print("lax.cond (proper dynamic branching):")
print(f"  jax_cond(5.0) = {jax_cond(5.0)}")
print(f"  jax_cond(-5.0) = {jax_cond(-5.0)}")
print("✓ Correct behavior!")

lax.cond (proper dynamic branching):
  jax_cond(5.0) = 5.0
  jax_cond(-5.0) = 5.0
✓ Correct behavior!

# BAD: Python for loop with traced loop bound
@jit
def python_loop(x, n):
    for i in range(n):  # range() needs concrete int!
        x = x + 1
    return x

print("Python for loop with traced n:")
try:
    result = python_loop(0.0, 3)
    print(f"  Loop with n=3: {result}")
except Exception as e:
    print(f"  ❌ Error: {type(e).__name__}")
    print("  range() requires a concrete integer, not a traced value!")

# WORKAROUND: Use static_argnums to make n concrete
from functools import partial

@partial(jit, static_argnums=(1,))
def python_loop_static(x, n):
    for i in range(n):  # Now n is concrete (static)
        x = x + 1
    return x

print("\nWith static_argnums (n is concrete):")
print(f"  Loop with n=3: {python_loop_static(0.0, 3)}")
print(f"  Loop with n=5: {python_loop_static(0.0, 5)}")
print("⚠️ Each different n causes recompilation!")

Python for loop with traced n:
  ❌ Error: TracerIntegerConversionError
  range() requires a concrete integer, not a traced value!

With static_argnums (n is concrete):
  Loop with n=3: 3.0

  Loop with n=5: 5.0
⚠️ Each different n causes recompilation!

# GOOD: Use lax.fori_loop for dynamic loops
@jit
def jax_loop(x, n):
    return lax.fori_loop(
        0, n,
        lambda i, x: x + 1,
        x
    )

print("lax.fori_loop (proper dynamic loop):")
print(f"  Loop with n=3: {jax_loop(0.0, 3)}")
print(f"  Loop with n=5: {jax_loop(0.0, 5)}")
print("✓ Single compiled function handles any n!")

lax.fori_loop (proper dynamic loop):

  Loop with n=3: 3.0
  Loop with n=5: 5.0
✓ Single compiled function handles any n!

5. Tracing and Concrete Values

Inside JIT, array values are abstract (only shape/dtype known). Operations requiring concrete values will fail.

# BAD: Using array value for Python control flow
@jit
def bad_concrete(x):
    if x[0] > 0:  # Needs concrete value!
        return x
    return -x

try:
    bad_concrete(jnp.array([1.0, 2.0]))
except jax.errors.TracerBoolConversionError as e:
    print(f"❌ Error: {type(e).__name__}")
    print("   Can't convert traced array to bool!")

❌ Error: TracerBoolConversionError
   Can't convert traced array to bool!

# GOOD: Use lax.cond or jnp.where
@jit
def good_concrete(x):
    # jnp.where works element-wise without needing concrete values
    return jnp.where(x > 0, x, -x)

print(f"good_concrete([1, -2, 3]) = {good_concrete(jnp.array([1.0, -2.0, 3.0]))}")
print("✓ Works correctly!")

good_concrete([1, -2, 3]) = [1. 2. 3.]
✓ Works correctly!

# Using static_argnums for values that should be concrete
from functools import partial

@partial(jit, static_argnums=(1,))  # Second arg (mode) is static
def with_static_arg(x, mode):
    if mode == "square":
        return x ** 2
    elif mode == "cube":
        return x ** 3
    else:
        return x

print("Using static_argnums:")
print(f"  mode='square': {with_static_arg(3.0, 'square')}")
print(f"  mode='cube': {with_static_arg(3.0, 'cube')}")
print("✓ Python control flow works with static arguments!")

Using static_argnums:

  mode='square': 9.0
  mode='cube': 27.0
✓ Python control flow works with static arguments!

6. Floating Point Precision

By default, JAX uses 32-bit floats! This can cause precision issues.

# Check default precision
x = jnp.array([1.0])
print(f"Default dtype: {x.dtype}")

# Precision issue example
a = jnp.array([1e-8])
b = jnp.array([1.0])

# In 32-bit, 1.0 + 1e-8 may equal 1.0!
result_32 = a + b
print(f"\n32-bit: 1.0 + 1e-8 = {result_32[0]}")
print(f"1e-8 was lost: {result_32[0] == 1.0}")

Default dtype: float32

32-bit: 1.0 + 1e-8 = 1.0
1e-8 was lost: True

# Enable 64-bit precision
jax.config.update("jax_enable_x64", True)

x = jnp.array([1.0])
print(f"With x64 enabled: {x.dtype}")

a = jnp.array([1e-8])
b = jnp.array([1.0])
result_64 = a + b

print(f"64-bit: 1.0 + 1e-8 = {result_64[0]:.16f}")
print(f"1e-8 preserved: {result_64[0] != 1.0}")

With x64 enabled: float64

64-bit: 1.0 + 1e-8 = 1.0000000099999999
1e-8 preserved: True

# Or specify dtype explicitly
x_32 = jnp.array([1.0], dtype=jnp.float32)
x_64 = jnp.array([1.0], dtype=jnp.float64)

print(f"Explicit float32: {x_32.dtype}")
print(f"Explicit float64: {x_64.dtype}")

Explicit float32: float32
Explicit float64: float64

7. NaN Debugging

NaNs can silently propagate through computations. JAX provides tools to catch them.

# NaN propagation example
def problematic_function(x):
    # sqrt of negative number = NaN
    return jnp.sqrt(x - 10) + x

x = jnp.array([5.0])  # 5 - 10 = -5, sqrt(-5) = NaN
result = problematic_function(x)
print(f"Result: {result}")
print("NaN silently appeared!")

Result: [nan]
NaN silently appeared!

# Enable NaN checking
from jax import config
config.update("jax_debug_nans", True)

try:
    result = problematic_function(x)
except FloatingPointError as e:
    print(f"Caught NaN: {e}")

# Disable for rest of notebook
config.update("jax_debug_nans", False)

Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.

Caught NaN: invalid value (nan) encountered in sqrt

# Manual NaN checking
def safe_function(x):
    # Check for valid input
    safe_input = jnp.maximum(x - 10, 0)  # Clamp negative values
    return jnp.sqrt(safe_input) + x

result = safe_function(x)
print(f"Safe result: {result}")
print(f"No NaN: {not jnp.isnan(result).any()}")

Safe result: [5.]

No NaN: True

8. Memory and Recompilation

JIT compilation has overhead. Understanding when recompilation happens is important.

# Recompilation due to shape changes
@jit
def process(x):
    return x ** 2 + x

print("Each new shape triggers recompilation:")

import time

for size in [10, 100, 1000, 10, 100]:
    x = jnp.ones(size)
    start = time.time()
    _ = process(x).block_until_ready()
    elapsed = time.time() - start
    print(f"  Size {size:4d}: {elapsed*1000:.2f}ms")

print("\nNotice: First call to each size is slow (compilation)")
print("Second call to same size is fast (cached)")

Each new shape triggers recompilation:

  Size   10: 48.90ms

  Size  100: 58.37ms

  Size 1000: 56.18ms
  Size   10: 0.04ms
  Size  100: 0.04ms

Notice: First call to each size is slow (compilation)
Second call to same size is fast (cached)

# Avoid recompilation with padding
def pad_to_max(x, max_size):
    """Pad array to fixed size."""
    padded = jnp.zeros(max_size)
    return padded.at[:len(x)].set(x), len(x)

@jit
def process_padded(x, length):
    # Only use first 'length' elements
    result = x ** 2 + x
    return result[:length]  # This still works but output shape varies

# Better: always return same shape
@jit 
def process_masked(x, mask):
    result = x ** 2 + x
    return jnp.where(mask, result, 0)

9. Pytree Gotchas

JAX’s pytree handling can be surprising.

# Pytrees flatten nested structures
params = {
    'layer1': {'W': jnp.ones((2, 3)), 'b': jnp.zeros(2)},
    'layer2': {'W': jnp.ones((4, 2)), 'b': jnp.zeros(4)},
}

# jax.tree.map applies function to all leaves
doubled = jax.tree.map(lambda x: x * 2, params)
print("Doubled params:")
print(f"  layer1.W sum: {doubled['layer1']['W'].sum()}")
print(f"  layer1.b sum: {doubled['layer1']['b'].sum()}")

Doubled params:
  layer1.W sum: 12.0
  layer1.b sum: 0.0

# GOTCHA: Lists vs tuples
# Lists are treated as pytree nodes (iterated over)
# Tuples can be either nodes or leaves depending on context

def show_leaves(tree, name):
    leaves = jax.tree.leaves(tree)
    print(f"{name}: {len(leaves)} leaves")
    for i, leaf in enumerate(leaves):
        print(f"  Leaf {i}: type={type(leaf).__name__}, shape={jnp.array(leaf).shape if hasattr(leaf, '__len__') else 'scalar'}")

# List is a pytree node
list_tree = [jnp.array([1, 2]), jnp.array([3, 4, 5])]
show_leaves(list_tree, "List")

# Dict is a pytree node
dict_tree = {'a': jnp.array([1, 2]), 'b': jnp.array([3, 4, 5])}
show_leaves(dict_tree, "Dict")

List: 2 leaves
  Leaf 0: type=ArrayImpl, shape=(2,)
  Leaf 1: type=ArrayImpl, shape=(3,)
Dict: 2 leaves
  Leaf 0: type=ArrayImpl, shape=(2,)
  Leaf 1: type=ArrayImpl, shape=(3,)

# GOTCHA: Pytree structure must match in tree_map with multiple trees

tree1 = {'a': jnp.array([1, 2]), 'b': jnp.array([3, 4])}
tree2 = {'a': jnp.array([10, 20]), 'b': jnp.array([30, 40])}

# This works - same structure
result = jax.tree.map(lambda x, y: x + y, tree1, tree2)
print(f"Matching structures: {result}")

# This fails - different structures
tree3 = {'a': jnp.array([10, 20]), 'c': jnp.array([30, 40])}  # Note: 'c' not 'b'
try:
    jax.tree.map(lambda x, y: x + y, tree1, tree3)
except ValueError as e:
    print(f"\n❌ Mismatched structures: {e}")

Matching structures: {'a': Array([11, 22], dtype=int64), 'b': Array([33, 44], dtype=int64)}

❌ Mismatched structures: Dict key mismatch; expected keys: ['a', 'b']; present keys: ['a', 'c'].

10. Common Error Messages

Here are common JAX errors and what they mean.

# Error: "Shapes must be 1D sequences of concrete values"
# Cause: Using traced value for shape

@jit
def bad_reshape(x, new_rows):
    return x.reshape((new_rows, -1))  # new_rows is traced!

try:
    bad_reshape(jnp.ones(6), 2)  # Want to reshape (6,) into (2, 3)
except Exception as e:
    print(f"Error type: {type(e).__name__}")
    print("Message: Shapes must use concrete values")

# Fix: Use static_argnums
from functools import partial

@partial(jit, static_argnums=(1,))
def good_reshape(x, new_rows):
    return x.reshape((new_rows, -1))  # Now new_rows is concrete

print(f"\n✓ Fixed: {good_reshape(jnp.ones(6), 2)}")  # (6,) -> (2, 3)

Error type: TypeError
Message: Shapes must use concrete values

✓ Fixed: [[1. 1. 1.]
 [1. 1. 1.]]

# Error: "grad requires real-valued outputs"
# Cause: Trying to differentiate w.r.t. non-scalar or complex output

def vector_output(x):
    return jnp.array([x, x**2, x**3])  # Returns vector!

try:
    grad(vector_output)(2.0)
except Exception as e:
    print(f"Error: {type(e).__name__}")
    print("grad needs scalar output!")

# Fix: Use jacfwd or jacrev for vector outputs
from jax import jacfwd
jacobian = jacfwd(vector_output)(2.0)
print(f"\n✓ Jacobian: {jacobian}")

Error: TypeError
grad needs scalar output!

✓ Jacobian: [ 1.  4. 12.]

# Error: "Tracer leaked"
# Cause: Storing traced values in global/external state

cache = {}

@jit
def bad_caching(x):
    cache['last'] = x  # Storing tracer in global dict!
    return x ** 2

# This might work but leads to subtle bugs
result = bad_caching(jnp.array([1.0, 2.0]))
print(f"Result: {result}")
print(f"Cache contains traced value: {type(cache.get('last'))}")
print("\n⚠️ Don't store traced values outside JIT!")

Result: [1. 4.]
Cache contains traced value: <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>

⚠️ Don't store traced values outside JIT!

Summary: Best Practices

Do:

• ✓ Write pure functions (no side effects, no global state)

• ✓ Use .at[].set() for array updates

• ✓ Split random keys before each use

• ✓ Use lax.cond, lax.fori_loop, lax.scan for control flow in JIT

• ✓ Enable jax_enable_x64 for numerical precision

• ✓ Use jax_debug_nans when debugging

• ✓ Batch operations with vmap

Don’t:

• ✗ Mutate arrays in-place

• ✗ Reuse random keys

• ✗ Use Python control flow for value-dependent branching in JIT

• ✗ Store traced values in global state

• ✗ Assume 64-bit precision (default is 32-bit)

• ✗ Create arrays with varying shapes inside JIT (causes recompilation)

# Quick reference: Control flow equivalents

print("""
Python → JAX Control Flow:
─────────────────────────────────────
if/else      → lax.cond(pred, true_fn, false_fn, operand)
for i in range(n)  → lax.fori_loop(0, n, body_fn, init_val)
while cond   → lax.while_loop(cond_fn, body_fn, init_val)
x[i] = v     → x.at[i].set(v)
x[i] += v    → x.at[i].add(v)

Random numbers:
─────────────────────────────────────
key = random.PRNGKey(seed)
key, subkey = random.split(key)
sample = random.normal(subkey, shape)

Debugging:
─────────────────────────────────────
jax.config.update("jax_debug_nans", True)
jax.config.update("jax_enable_x64", True)
jax.debug.print("{x}", x=value)  # Print inside JIT
""")

Python → JAX Control Flow:
─────────────────────────────────────
if/else      → lax.cond(pred, true_fn, false_fn, operand)
for i in range(n)  → lax.fori_loop(0, n, body_fn, init_val)
while cond   → lax.while_loop(cond_fn, body_fn, init_val)
x[i] = v     → x.at[i].set(v)
x[i] += v    → x.at[i].add(v)

Random numbers:
─────────────────────────────────────
key = random.PRNGKey(seed)
key, subkey = random.split(key)
sample = random.normal(subkey, shape)

Debugging:
─────────────────────────────────────
jax.config.update("jax_debug_nans", True)
jax.config.update("jax_enable_x64", True)
jax.debug.print("{x}", x=value)  # Print inside JIT