Skip to content

Chapter 3

Overview

JAX requires functional programming. This chapter covers pure functions and immutable data structures essential for JAX.

-

What is a Pure Function?

A pure function satisfies:

  1. Deterministic: Same input → always same output
  2. No side effects: Doesn't modify external state
  3. No I/O: Doesn't read/write files or network
  4. Referentially transparent: Can be replaced with return value

Example: Pure vs Impure

# IMPURE
total = 0

def add_impure(x):
 global total
 total += x
 print(f"Total: {total}") # I/O side effect
 return total

result1 = add_impure(5)
result2 = add_impure(5)
# Same input, different outputs (5 vs 10)!

# PURE
def add_pure(x):
 return x + x

result1 = add_pure(5) # 10
result2 = add_pure(5) # 10 (always same)

Why Pure Functions Matter in JAX

The Tracing Problem

When you call jax.grad(f), JAX traces your function:

import jax
import jax.numpy as jnp

def f_pure(x):
 return x**2 + jnp.sin(x)

# JAX traces through f with symbolic values:
# 1. Executes with symbolic x
# 2. Records operations
# 3. Computes derivative of recorded graph
grad_f = jax.grad(f_pure)
print(grad_f(3.0)) # Works: 2*3 + cos(3) ≈ 6.96

With side effects, tracing breaks:

import jax
import jax.numpy as jnp

global_counter = 0

def f_impure(x):
 global global_counter
 global_counter += 1 # Side effect: modifies external state
 return x**2

# During tracing:
# JAX increments global_counter, but doesn't record it
# When computing gradient, counter state is inconsistent
grad_f = jax.grad(f_impure)
print(grad_f(3.0)) # Unpredictable result!

Composition Requires Purity

JAX transformations compose only when functions are pure:

import jax
import jax.numpy as jnp

# Pure function
def f(x):
 return x**2

# Compose transformations
gradient_of_vmap = jax.grad(jax.vmap(f))
# = Gradient of (Apply f to batch)
# = Works correctly for pure functions

# With impure function:
global_state = []

def g_impure(x):
 global_state.append(x) # Side effect
 return x**2

# Composing transformations on impure functions:
try:
 jax.grad(jax.vmap(g_impure)) # Unpredictable behavior!
except:
 pass

Eliminating Side Effects

Pattern 1: Return Results Instead of Modifying

# IMPURE
def append_impure(lst, item):
 lst.append(item) # Modifies original!
 return lst

my_list = [1, 2, 3]
result = append_impure(my_list, 4)
print(my_list) # [1, 2, 3, 4] - MODIFIED!

# PURE
def append_pure(lst, item):
 return lst + [item] # New list

my_list = [1, 2, 3]
result = append_pure(my_list, 4)
print(my_list) # [1, 2, 3] - unchanged
print(result) # [1, 2, 3, 4]

Pattern 2: Pass Dependencies as Arguments

# IMPURE
config = {"learning_rate": 0.01}

def train_step_impure(weights):
 lr = config["learning_rate"] # Depends on global state
 return weights - lr * gradients

# PURE
def train_step_pure(weights, learning_rate):
 return weights - learning_rate * gradients

# Now reusable with different learning rates
weights1 = train_step_pure(weights, 0.01)
weights2 = train_step_pure(weights, 0.001)

Pattern 3: Return Multiple Values

# IMPURE
class State:
 def __init__(self):
 self.count = 0
 self.values = []

state = State()

def process_impure(x):
 state.count += 1
 state.values.append(x)
 return x**2

# PURE
def process_pure(state, x):
 new_count = state['count'] + 1
 new_values = state['values'] + [x]
 new_state = {'count': new_count, 'values': new_values}
 result = x**2
 return result, new_state

state = {'count': 0, 'values': []}
result, state = process_pure(state, 5)

-

Immutability in JAX

JAX Arrays are Immutable

import jax.numpy as jnp

x = jnp.array([1, 2, 3])

# Can't modify in place
# x[0] = 999 # TypeError

# Create new array with.at[]
x = x.at[0].set(999)
print(x) # [999 2 3]

Immutable Update Methods

import jax.numpy as jnp

x = jnp.array([1, 2, 3, 4, 5])

# set
x = x.at[0].set(999)
print(x) # [999 2 3 4 5]

# add
x = jnp.array([1, 2, 3])
x = x.at[0].add(10)
print(x) # [11 2 3]

# multiply
x = jnp.array([1, 2, 3])
x = x.at[1].multiply(100)
print(x) # [1 200 3]

# subtract, divide, min, max
x = x.at[2].subtract(5) # x[2] -= 5
x = x.at[2].divide(2) # x[2] /= 2

# Batch updates
x = jnp.array([1, 2, 3, 4, 5])
x = x.at[0, 2, 4](/0,-2,-4/).set(999)
print(x) # [999 2 999 4 999]

Chaining Updates

import jax.numpy as jnp

x = jnp.array([1, 2, 3, 4, 5])

# Chain multiple updates
x = (x.at[0].add(10)
.at[2].multiply(100)
.at[4].set(999))

print(x) # [11 2 300 4 999]

Deep Immutability

import jax.numpy as jnp

# Nested structures
data = {
 'A': jnp.array([1, 2, 3]),
 'B': jnp.array([4, 5, 6])
}

# Update array in dict
data['A'] = data['A'].at[0].set(999)

# Alternative
data = {
 **data,
 'A': data['A'].at[0].set(999)
}

Functional Programming Patterns

Higher-Order Functions

Functions that return functions:

# Return a configured function
def make_adder(n):
 def add(x):
 return x + n
 return add

add_5 = make_adder(5)
print(add_5(10)) # 15
print(add_5(20)) # 25

Closures

Outer function's variables captured:

def make_multiplier(factor):
 def multiply(x):
 return x * factor
 return multiply

times_3 = make_multiplier(3)
times_5 = make_multiplier(5)

print(times_3(10)) # 30
print(times_5(10)) # 50

Function Composition

Combining functions:

def compose(f, g):
 """Returns h(x) = f(g(x))"""
 def h(x):
 return f(g(x))
 return h

# Example
def double(x):
 return x * 2

def add_one(x):
 return x + 1

double_then_add = compose(add_one, double)
print(double_then_add(5)) # add_one(double(5)) = 11

-

Handling State Functionally

State as Arguments/Returns

import jax.numpy as jnp

# Traditional
# class Accumulator:
# def __init__(self):
# self.total = 0
# def add(self, x):
# self.total += x # Impure!

# Functional
def accumulate_step(state, x):
 """Returns (result, new_state)"""
 new_total = state['total'] + x
 new_state = {**state, 'total': new_total}
 return new_total, new_state

# Usage
state = {'total': 0}
result1, state = accumulate_step(state, 5)
result2, state = accumulate_step(state, 3)
result3, state = accumulate_step(state, 7)
print(state) # {'total': 15}

Training Loop Example

import jax
import jax.numpy as jnp
from jax import grad

def loss(params, x, y):
 pred = params['w'] * x + params['b']
 return jnp.mean((pred - y)**2)

def train_step(params, x, y, learning_rate):
 """Pure: takes state, returns new state"""
 grads = grad(loss)(params, x, y)
 new_params = {
 'w': params['w'] - learning_rate * grads['w'],
 'b': params['b'] - learning_rate * grads['b']
 }
 loss_value = loss(new_params, x, y)
 return new_params, loss_value

# Usage
params = {'w': 0.0, 'b': 0.0}
x, y = jnp.array([1., 2., 3.]), jnp.array([2., 4., 6.])

for step in range(100):
 params, loss_val = train_step(params, x, y, 0.01)

print(params) # Learned parameters

JAX-Specific Patterns

Using lax for Control Flow

import jax
import jax.lax as lax
import jax.numpy as jnp

# Pure conditional
def absolute_value(x):
 return lax.cond(
 x < 0,
 lambda: -x, # True branch
 lambda: x # False branch
)

print(absolute_value(-5)) # 5
print(absolute_value(5)) # 5

# Pure loop (scan)
def sum_n(n):
 def body(carry, x):
 return carry + x, None

 result, _ = lax.scan(body, 0, jnp.arange(n))
 return result

print(sum_n(10)) # 45 (sum of 0 to 9)

Using JAX Pytrees

JAX can work with nested structures automatically:

import jax
import jax.numpy as jnp

# Nested structure (pytree)
tree = {
 'a': jnp.array([1., 2., 3.]),
 'b': jnp.array([4., 5.]),
 'c': [jnp.array([6.]), jnp.array([7., 8.])]
}

# JAX operations work on trees
def add_one_to_all(tree):
 # JAX automatically maps over structure
 return jax.tree_map(lambda x: x + 1, tree)

result = add_one_to_all(tree)
print(result['a']) # [2. 3. 4.]

# Flatten tree for operations
leaves, treedef = jax.tree_util.tree_flatten(tree)
print(len(leaves)) # 4 (flattened arrays)

# Reconstruct
tree_reconstructed = jax.tree_util.tree_unflatten(treedef, leaves)

Common Mistakes

Mistake 1: Assuming Mutability

import jax.numpy as jnp

x = jnp.array([1, 2, 3])

# This doesn't modify x
y = x.reshape((1, 3))
# x is still shape (3,)

# This doesn't modify x
z = x + 5
# x is still [1, 2, 3]

# Reassign to modify "state"
x = x.reshape((1, 3))
x = x + 5

Mistake 2: Modifying Function Arguments

# IMPURE
def process_list_bad(lst):
 lst.append(999) # Modifies input!
 return sum(lst)

# PURE
def process_list_good(lst):
 new_lst = lst + [999] # New list
 return sum(new_lst)

Mistake 3: Using Global Variables

# IMPURE
global_config = {'lr': 0.01}

def train_step(params):
 return params - global_config['lr'] * grads

# PURE
def train_step(params, lr):
 return params - lr * grads

-

Best Practices

1. Design for Immutability

# Prefer functional approach
def process_data(data, config):
 # Don't modify data or config
 processed = transform(data)
 new_config = {**config, 'updated': True}
 return processed, new_config

2. Use Return Values for State

# Return modified state
state, loss = train_step(state, batch)
state, new_loss = train_step(state, next_batch)

3. Pass Dependencies as Arguments

# All needed information as arguments
def compute(x, y, z, config):
 # Don't rely on globals
 return f(x, y, z, config['param'])

4. Use JAX Utilities for Control Flow

# Use lax.cond, lax.scan, lax.while_loop
# Not Python if/for statements

-

Summary

PURE FUNCTION CHECKLIST:
Same input → same output (deterministic)
No global state modifications
No I/O (files, network, print in grad)
No randomness (use jax.random properly)
Returns all output values
Takes all inputs as arguments

-

Next Steps

Checkpoint: Explain why JAX requires pure functions (hint: tracing and composition)

-

Last Updated: 2026-08-09