What Breaks Purity: The Rogues Gallery

Let's catalog every common way purity gets violated in JAX, so you can spot these patterns immediately.

1. In-Place Mutation

import jax.numpy as jnp

x = jnp.array([1, 2, 3])
# x[0] = 99  # TypeError: JAX arrays are immutable

# Fix: use .at[].set()
x_new = x.at[0].set(99)  # Returns new array, x is unchanged

2. Global / Closure Variables

import jax
import jax.numpy as jnp

learning_rate = 0.01

# BAD: reads from closure
@jax.jit
def update_bad(params, grads):
    return params - learning_rate * grads

# GOOD: pass as argument
@jax.jit
def update_good(params, grads, lr):
    return params - lr * grads

# OR: use static_argnums for values that rarely change
@jax.jit
def update_static(params, grads, lr):
    return params - lr * grads
# JAX will recompile when lr changes, but that's acceptable if it rarely does

3. Python print() Under JIT

import jax
import jax.numpy as jnp

@jax.jit
def fn_with_print(x):
    print("This runs during TRACING only, not execution!")
    y = x + 1
    print(f"y = {y}")  # Prints a tracer object, not a number
    return y

result = fn_with_print(jnp.array(5.0))
# Output during first call:
# "This runs during TRACING only, not execution!"
# "y = Traced<ShapedArray(float32[])>with<DynamicJaxprTrace...>"

# Second call: no print at all — JIT reuses the cached trace
result2 = fn_with_print(jnp.array(10.0))

For debugging, use jax.debug.print instead:

@jax.jit
def fn_with_debug_print(x):
    y = x + 1
    jax.debug.print("y = {}", y)  # Prints at execution time!
    return y

fn_with_debug_print(jnp.array(5.0))   # Prints "y = 6.0"
fn_with_debug_print(jnp.array(10.0))  # Prints "y = 11.0"

4. Random Number Generation

import jax
import jax.numpy as jnp

# NumPy uses global state — IMPURE
import numpy as np
np.random.seed(42)
a = np.random.randn(3)  # Mutates global RNG state
b = np.random.randn(3)  # Different result — depends on hidden state

# JAX uses explicit keys — PURE
key = jax.random.PRNGKey(42)
key1, key2 = jax.random.split(key)
a = jax.random.normal(key1, (3,))  # Deterministic given key1
b = jax.random.normal(key2, (3,))  # Deterministic given key2

# Same key always gives same result
a_again = jax.random.normal(key1, (3,))
print(jnp.allclose(a, a_again))  # True — pure!

5. Python Control Flow That Depends on Values

import jax
import jax.numpy as jnp

# PROBLEMATIC under JIT: Python if depends on a traced value
@jax.jit
def bad_relu(x):
    if x > 0:  # ConcretizationTypeError!
        return x
    else:
        return 0.0

# GOOD: use JAX control flow
@jax.jit
def good_relu(x):
    return jnp.where(x > 0, x, 0.0)

We'll cover JIT control flow in detail in Track 4.