A side effect changes state outside a function’s explicit return value. In JAX, hidden variables, global state, printing, and other side effects conflict with the pure-function style required by transformations such as JIT compilation. section 1.1.2 section 2.9
Pure JAX functions derive behavior only from their inputs: identical inputs produce identical outputs. Parameters, optimizer state, gradients, and random-number-generator state therefore travel explicitly through function arguments rather than hidden object state. section 2.9 section 2.2
Avoid side effects by making state and updates explicit:
- Pass model parameters, optimizer state, gradients, and random-number-generator state into functions. section 1.1.2 section 2.9
- Treat JAX arrays as immutable; replace NumPy-style in-place updates such as x[idx] = y with x = x.at[idx].set(y). section 3.3.1
- Keep global state and observable actions such as print statements out of functions transformed with JIT. section 5.1.2
A side effect can disappear after JIT tracing
A function that prints x and multiplies by global_state may print during its first jitted call, when JAX traces and compiles it, but not during later calls. Changing global_state afterward does not change the cached compiled behavior. The returned value can therefore reflect the first traced execution rather than the function’s later Python behavior. section 5.1.2
Listing 5.9 Compiling an impure function
global_state = 1 #1
def impure_function(x):
print(f'Side-effect: printing x={x}') #2
y = x*global_state #3
return y
impure_function_jit = jax.jit(impure_function) #4
impure_function_jit(10)
>>> Side-effect: printing x=Traced<ShapedArray(int32[], #5
weak_type=True)>with<DynamicJaxprTrace(level=1/0)> #5
>>> Array(10, dtype=int32, weak_type=True) #5
impure_function_jit(10)
>>> Array(10, dtype=int32, weak_type=True) #6
global_state = 2 #7
impure_function_jit(10) #8
>>> Array(10, dtype=int32, weak_type=True)
impure_function(10) #9
>>> Side-effect: printing x=10
>>> 20
Inspect the runnable impure-function example showing a print side effect, global state, first-call tracing, and later cached calls. section 5.1.2
