What Pure Means and Why JAX Demands It

A pure function is a function where the output depends only on the inputs, and the function has no side effects. This is the single most important concept in JAX programming, and violating it is the source of most JAX bugs.

Two rules define purity:

1. Deterministic: Same inputs always produce the same outputs.
2. No side effects: The function doesn't modify anything outside itself — no global variables, no mutating arguments, no I/O.

import jax.numpy as jnp

# PURE: output depends only on input, no side effects
def pure_fn(x):
    return jnp.sum(x ** 2)

# IMPURE: depends on global state
scale = 2.0
def impure_global(x):
    return jnp.sum(x ** 2) * scale  # Reads global variable!

# IMPURE: has side effects
results = []
def impure_sideeffect(x):
    result = jnp.sum(x ** 2)
    results.append(result)  # Side effect: modifies external list!
    return result

# IMPURE: mutates input
def impure_mutation(x):
    x[0] = 0  # Side effect: modifies input! (also fails in JAX)
    return jnp.sum(x)

Why Does JAX Care About Purity?

JAX's transformations — jit, grad, vmap, pmap — all work by tracing your function. During tracing, JAX runs your function with abstract placeholder values (called "tracers") instead of real data. It records what operations happen and builds a computation graph.

This only works if the function is pure:

• jit traces your function once, then replays the recorded computation. If your function reads global state, the traced version will use whatever the global value was at trace time, ignoring future changes.
• grad needs to know the exact mathematical function to differentiate. Side effects aren't differentiable.
• vmap needs to know that running the function on different inputs produces independent results. Side effects (like appending to a list) could create dependencies between "parallel" executions.

import jax
import jax.numpy as jnp

# Demonstration: global state is captured at trace time
multiplier = 2.0

@jax.jit
def buggy_multiply(x):
    return x * multiplier

print(buggy_multiply(jnp.array(3.0)))  # 6.0

multiplier = 10.0  # Change the global
print(buggy_multiply(jnp.array(3.0)))  # Still 6.0! JIT cached the old value

The correct approach:

@jax.jit
def correct_multiply(x, multiplier):
    return x * multiplier

print(correct_multiply(jnp.array(3.0), 2.0))   # 6.0
print(correct_multiply(jnp.array(3.0), 10.0))  # 30.0 — correct!