Well be back guys, in our previous article on managing PRNG keys, we explored how JAX abandons the conventional hidden random state in favor of stateless random number generation. We learned that to get deterministic sequences of random numbers, we must pass keys into functions and split them as state updates.
We now zoom out and tackle a broader design pattern that governs all execution in JAX as stated in topic.
For those familiar with Object-oriented ML frameworks like PyTorch or Tensorflow, the mental model of state likely revolves around object attributes that get modified in-place during computation. In JAX, functional purity means that hidden mutations are fundamentally incompatible with compilation with jax.jit and automatic differentiation with jax.grad
In this article, we'll cover how state works in pure functional programming, how to express iterations and conditions using jax.lax primitives and hopefully build a training loop that handles parameters, optimizer state and PRNG keys in a unified state.
Why state must be handled functionally
In traditional code, functions often carry side effects. They may read or mutate global variables, alter internal class properties or perform inplace updates like x+=1. Consider the sample code below
class Counter:
def __init__(self):
self.count=0
def increment(self, amount):
self.count += amount # This mutates the internal state
return self.count
As expressed in earlier articles about how JAX transformations like jax.jit and jax.grad work using tracers, if a function relies on side effects or internal mutations, the trace will fail to capture the changes that happen outside the return values or hardcoded values during a single trace.
To avoid all that, JAX requires pure functions that if given the same input, they always produce the same output and above all have zero side effects on state or mutation.
The Explicit State / Carry Pattern
To update parameters, track stats or counters in JAX, we use the Carry Pattern. Instead of updating an object in-place, your function accepts the state as an argument(the carry), performs the required updates using immutable operations and returns the newly computed state alongside any results.
Consider the Counter example code below
import jax
import jax.numpy as jnp
# State management using a pure function
def increment(state: jnp.ndarray, amount: float) -> tuple[jnp.ndarray, jnp.ndarray]
"""This pure function takes the current state, computes the next state and returns both the new state and output """
new_state = state + amount
return new_state, new_state
initial_state = jnp.array(0.0)
# Step 1
state, output1 = increment(initial_state, 1.0)
print(f"Output: {output1}, New State: {state}")
# Step 2
new_state, output2 = increment(state, 2.5)
print(f"Output: {output2}, New State: {new_state}")
Output
Output: 1.0, New State: 1.0
Output: 3.5, New State: 3.5
By making state flow through arguments and return values, JAX compilers can optimize operations, parallelize computation across hardware targets and safely differentiate across multi-step algorithms.
Low-Level Control flow primitives with jax.lax
Writing Python conventional flow statements like for loops or if/else statements inside jax.jit compiled functions presents a unique set of challenges. Forexample, a standard for loop with 1000 iterations creates a Jaxpr computational graph with 1000 repeated blocks leading to high compilation overhead. While standard if statements evaluated on tracer values trigger a ConcretizationTypeError because JAX cannot determine truth values of abstract shapes during tracing. To write compiled loops and conditionals that scale efficiently, JAX provides primitives inside the low-level jax.lax module.
Sequential Iteration with jax.lax.scan
jax.lax.scan is JAX's workhorse primitive for sequential and recurrent processing. It executes a loop where an accumulated state, the carry, is passed from step to step while optionally consuming a sequence of inputs and outputting a sequence of results.
Syntax
final_carry, stacked_outputs = jax.lax.scan(
f,
init,
xs,
Length=None
)
where
f is the expected step function f(carry, x) -> (next_carry, y) where the carry is the currently accumulated state and x is the single slice of the current iteration step xs[i] and returns a two-element tuple next_carry which is the updated step i+1 and y, the perstep output you may save.
init is the starting carry value before the first iteration begins. This can be a scalar, JAX Array or a PyTree.
xs being the sequence of inputs to iterate through.
The return values are the final_carry which contains the updated state after processing the last element of xs with the same structure and shape as init and stacked_outputs containing all per-step outputs(y) concatenated together along axis 0 with a shape (loop_length, *shape_of_y)
Consider running a cumulative sum
Let's use lax.scan to compute a running sum across an array while tracking the current aggregate state.
import jax
import jax.numpy as jnp
# step fn
def scan_running_sum(carry: jnp.ndarray, x: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
"""
carry: running total so far
x: individual element from input sequence
"""
new_carry = carry + x
output = new_carry
return new_carry, output
xs = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
init_state = jnp.array(0.0)
final_carry, running_totals = jax.lax.scan(scan_running_sum, init_state, xs)
print(f"Inputs: {xs}")
print(f"Final carry: {final_carry}")
print(f"Running Totals: {running_totals}")
Output
Inputs: [1. 2. 3. 4. 5.]
Final carry: 15.0
Running Totals: [ 1. 3. 6. 10. 15.]
Consider a Recurrent Neural Network (RNN) Step
lax.scan is ideal for sequential models like Recurrent Neural Networks (RNNs) or Transformers processing temporal sequences because it compiles the iteration into a single optimized loop rather than thousands of unrolled nodes. Take an example of the codebase below
import jax
import jax.numpy as jnp
def rnn_cell(hidden_state: jnp.ndarray, x_t: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
"""
step: h_t = tanh(W_h * h_{t-1} + W_x * x_t)
fixed identity-like weights are used for demo purposes
"""
W_h = jnp.array([[0.8, 0.1], [0.0, 0.8]])
W_x = jnp.array([[1.0, 0.0], [0.0, 1.0]])
next_hidden = jnp.tanh(jnp.dot(hidden_state, W_h) + jnp.dot(x_t, W_x))
return next_hidden, next_hidden
sequence_inputs = jnp.array([
[1.0, 0.5],
[0.2, 0.8],
[0.0, 1.0],
[-0.5, 0.0]
])
init_hidden = jnp.zeros((2,))
# Trace and execute sequence via scan
final_hidden, hidden_states_history = jax.lax.scan(rnn_cell, init_hidden, sequence_inputs)
print(f"Sequence Length: {sequence_inputs.shape[0]}")
print(f"Hidden States Shape over time: {hidden_states_history.shape}")
Output
Sequence Length: 4
Hidden States Shape over time: (4, 2)
Conditional Execution with jax.lax.cond
This is the JIT-compatible equivalent of an if-else statement. This is so because in a typical Python if-else, execution is dynamic based on values evaluated at runtime. Since JAX's jax.jit needs to build a static computational graph, jax.lad.cond allows JAX to trace both branches during compilation and execute only the appropriate branch based on a runtime boolean array.
Syntax
result = jax.lax.cond(
pred,
true_fn,
false_fn,
operand
)
where
pred is a boolean scalar array like jnp.array(True) or a conditional statement like x>0. It cannot be a multi-element array.
true_fn and false_fn are the signature alignments that must accept a single argument passed via the operand. These are the independent paths based on the evaluation of the pred. If the functions must take multiple arguments, pass them as a tuple. These must also return outputs with identical shapes and dtypes as required by the static compilation graph to keep a fixed shape.
operand is the data to be passed either to the true_fn or false_fn
Consider the gradient clipping example below
import jax
import jax.numpy as jnp
def threshold_activation(x: jnp.ndarray) -> jnp.ndarray:
"""Applies custom piece scaling"""
pred = jnp.mean(x) > 0.0
true_fn = lambda val: val * 2.0
false_fn = lambda val: val * 0.1
return jax.lax.cond(pred, true_fn, false_fn, x)
compiled_activation = jax.jit(threshold_activation)
pos_data = jnp.array([0.5, 1.5, 2.0])
neg_data = jnp.array([-1.0, -0.5, -0.2])
print("Positive batch result:", compiled_activation(pos_data))
print("Negative batch result:", compiled_activation(neg_data))
Output
Positive batch result: [1. 3. 4.]
Negative batch result: [-0.1 -0.05 -0.02]
Dynamic Iterations with jax.lax.while_loop
During compilation, a standard while loop gets unrolled and requires a fixed compile-time iteration count, while a jax.lax.while_loop allows you to run dynamic condition-based loops without unrolling the computation graph.
Syntax
final_val = jax.lax.while_loop(
cond_func,
body_func,
init_val
)
where
cond_func takes the current loop state value starting with the init_val and returns a boolean scalar array i.e. True to keep looping and False to stop.
body_func takes the current loop state value and returns an updated state of the exact same structure, shape and data type.
init_val is the starting state object usually a PyTree passed into a loop.
Consider the code example to find a convergence below
import jax
import jax.numpy as jnp
@jax.jit
def find_next_power_of_two(target):
# loop state carries
# (current_value, step_count)
init_val = (1, 0)
# checks if current_value < target
def cond_func(val):
current_val, count = val
return current_val < target
# doubles current_value and increments count
def body_func(val):
current_val, count = val
return (current_val * 2, count + 1)
final_val, total_steps = jax.lax.while_loop(cond_func, body_func, init_val)
return final_val, total_steps
print(find_next_power_of_two(50)) #(64, 6)
Output
(Array(64, dtype=int32, weak_type=True), Array(6, dtype=int32, weak_type=True))
Note: Because the loop count isn't fixed, reverse-mode automatic differentiation (jax.grad) cannot automatically differentiate through a while_loop without additional bounds.
The power of JAX reveals itself when combining these functional state principles with core transformations i.e jax.jit and jax.grad. Because state updates are represented explicitly as inputs and outputs, wrapping state-modifying functions with JAX transformations is seamless and free of hidden side effects.
import jax
import jax.numpy as jnp
from typing import Dict, Tuple, Any
def compute_loss(params: Dict[str, jnp.ndarray], x: jnp.ndarray, y: jnp.ndarray, key: jax.Array) -> jnp.ndarray:
""" Compute mean squared error with stochastic noise injection """
# Inject stochasticity using PRNG key
noise = jax.random.normal(key, shape=x.shape) * 0.01
x_noisy = x + noise
# model forward pass
predictions = jnp.dot(x_noisy, params['w']) + params['b']
loss = jnp.mean((predictions - y) ** 2)
return loss
def train_step(
state: Tuple[Dict[str, jnp.ndarray], Dict[str, jnp.ndarray], jax.Array],
batch: Tuple[jnp.ndarray, jnp.ndarray]
) -> Tuple[Tuple[Dict[str, jnp.ndarray], Dict[str, jnp.ndarray], jax.Array], jnp.ndarray]:
"""
Performs a single SGD update with momentum and updates the key.
Carry State: (params, opt_state, key)
Batch: (x, y)
"""
params, opt_state, key = state
x, y = batch
# Split the keys
key, subkey = jax.random.split(key)
# Compute loss and gradient with respect to params (arg 0)
loss, grads = jax.value_and_grad(compute_loss, argnums=0)(params, x, y, subkey)
# Momentum SGD Optimizer Logic (Pure functional updates)
learning_rate = 0.01
momentum = 0.9
new_opt_state = {}
new_params = {}
for param_name in params:
# Update momentum state: v_t = mu * v_{t-1} + g_t
velocity = momentum * opt_state[param_name] + grads[param_name]
new_opt_state[param_name] = velocity
# Update weights: w_t = w_{t-1} - lr * v_t
new_params[param_name] = params[param_name] - learning_rate * velocity
new_state = (new_params, new_opt_state, key)
return new_state, loss
#Initialize inputs, state, and run demo updates
rng = jax.random.key(42)
rng, subkey = jax.random.split(rng)
# Model state initialization
params = {
'w': jax.random.normal(subkey, (3, 1)),
'b': jnp.zeros((1,))
}
# Optimizer state
opt_state = {
'w': jnp.zeros_like(params['w']),
'b': jnp.zeros_like(params['b'])
}
# Bundle state into carry tuple
state = (params, opt_state, rng)
# Mock Batch Data
x_batch = jnp.ones((8, 3))
y_batch = jnp.ones((8, 1)) * 2.0
# Perform 3 updates
for step in range(3):
state, loss_val = train_step(state, (x_batch, y_batch))
print(f"Step {step + 1} | Loss: {loss_val:.4f}")
Output:
Step 1 | Loss: 2.2618
Step 2 | Loss: 1.9221
Step 3 | Loss: 1.3474
If you look closely at the signature and return value of the train step, you'll notice how everything the step needs to know about the current world is packed into a single tuple.
Inside the step, you'll also notice a couple of functional operations happening simultaneously.
We split the PRNG key so we can safely inject random noise into compute_loss using subkey, while saving key for the next iteration.
jax.value_and_grad calculates both the scalar loss and the exact parameter gradients with respect to params.
Instead of modifying weights in-place like param -= lr * grad, we compute brand-new dictionaries for new_opt_state and new_params.
Finally, we pack those updated values back into a new_state tuple and return it alongside the loss scalar.
When we run the outer for loop, we simply overwrite our local state variable with the updated tuple returned by train_step. Because train_step is decorated with @jax.jit, JAX compiles this entire pipeline into a single, highly optimized XLA executable that runs directly on your accelerator with zero side effects.
All in all
In JAX, pure functions require state to be explicitly passed in as an argument and explicitly returned as a transformed output (the Carry Pattern).
Replace Python for loops in sequential or recurrent workflows with lax.scan to compile long sequences into optimized single-trace instructions.
Use JAX-native control flow primitives whenever branching logic or termination conditions depend on dynamic runtime arrays.
Combining explicit carry patterns with transformations like jax.jit and jax.grad allows you to write ultra-fast, auto-differentiated training logic without state leakage or compiler errors.
See you guys in the next one.
