Let's take a realistic piece of impure numerical code and systematically refactor it into pure JAX. This is the exact process you'll follow when porting existing code.
The Impure Version
import numpy as np
# Global state — impure!
class RunningStats:
def __init__(self):
self.count = 0
self.mean = 0.0
self.variance = 0.0
def update(self, batch):
"""Welford's online algorithm — mutates self."""
for x in batch:
self.count += 1
delta = x - self.mean
self.mean += delta / self.count
delta2 = x - self.mean
self.variance += delta * delta2
def get_stats(self):
return self.mean, self.variance / max(self.count - 1, 1)
# Usage: side effects everywhere
stats = RunningStats()
for _ in range(10):
batch = np.random.randn(100)
stats.update(batch)
mean, var = stats.get_stats()
print(f"Mean: {mean:.4f}, Var: {var:.4f}")
The Pure JAX Version
import jax
import jax.numpy as jnp
# State is a plain data structure (NamedTuple or dict)
from typing import NamedTuple
class StatsState(NamedTuple):
count: jnp.ndarray
mean: jnp.ndarray
m2: jnp.ndarray # Sum of squared deviations
def init_stats():
"""Create initial state."""
return StatsState(
count=jnp.array(0, dtype=jnp.int32),
mean=jnp.array(0.0),
m2=jnp.array(0.0),
)
def update_stats(state, x):
"""Welford's algorithm — pure function, returns new state."""
count = state.count + 1
delta = x - state.mean
mean = state.mean + delta / count
delta2 = x - mean
m2 = state.m2 + delta * delta2
return StatsState(count=count, mean=mean, m2=m2)
def get_stats(state):
"""Extract mean and variance from state."""
variance = jnp.where(state.count > 1, state.m2 / (state.count - 1), 0.0)
return state.mean, variance
# Process a batch using jax.lax.scan (functional loop)
@jax.jit
def process_batch(state, batch):
state = jax.lax.fori_loop(
0, batch.shape[0],
lambda i, s: update_stats(s, batch[i]),
state
)
return state
# Usage: explicit state threading, no mutation
state = init_stats()
key = jax.random.PRNGKey(42)
for i in range(10):
key, subkey = jax.random.split(key)
batch = jax.random.normal(subkey, (100,))
state = process_batch(state, batch)
mean, var = get_stats(state)
print(f"Step {i}: Mean: {mean:.4f}, Var: {var:.4f}")
What Changed in the Refactoring
| Impure Pattern | Pure Pattern |
|---|---|
| Class with mutable attributes | NamedTuple (immutable data) |
Methods that modify self | Functions that return new state |
self.count += 1 | count = state.count + 1 |
| Python for loop | jax.lax.fori_loop |
np.random.randn() | jax.random.normal(key, ...) |
| Implicit state in object | Explicit state passed between functions |
💡 Why This Matters
The functional version is more verbose, but it's now fully compatible with all JAX transformations. You could jit it (done above), grad through it (if you needed gradients of the statistics), or vmap it across multiple independent streams. The explicit state also makes serialization trivial — just save the StatsState namedtuple.