The problem with untraceable state#

Let’s start with a simple stateful program: a counter.

import jax
import jax.numpy as jnp

class Counter:
  """A simple counter."""

  def __init__(self):
    self.n = 0

  def count(self) -> int:
    """Increments the counter and returns the new value."""
    self.n += 1
    return self.n

counter = Counter()

for _ in range(3):
  print(counter.count())

1
2
3

The n attribute maintains the counter’s state between calls, updated as a side effect. As plain Python this works fine, but it falls apart under JAX transformations. We can see exactly how by looking at what tracing records, using the jaxpr-inspection idiom from How transformations work: tracing:

counter = Counter()
jax.jit(counter.count).trace().jaxpr

{ lambda ; . let  in (1:i32[],) }

The recorded program takes no input, performs no operations, and always returns 1. The side effect self.n += 1 ran once, at trace time; the value self.n produced was captured as a constant; and the state update itself is nowhere in the jaxpr. Any transformation working from this traced view is working with a function that returns 1, forever. (Under jax.jit, which caches traces, that’s precisely what you’d observe: the compiled counter returns 1 on every call.)

Threading state through pure functions#

The fix is to make the state part of the function’s interface: take the current state as an argument, and return the updated state alongside the output.

CounterState = int

class CounterV2:

  def count(self, n: CounterState) -> tuple[int, CounterState]:
    # You could just return n+1 once, but here we separate its role as
    # the output and as the counter state for didactic purposes.
    return n + 1, n + 1

  def reset(self) -> CounterState:
    return 0

counter = CounterV2()
state = counter.reset()

for _ in range(3):
  value, state = counter.count(state)
  print(value)

1
2
3

The caller now keeps track of the state explicitly. In exchange, count is pure and hence easy to trace, and its jaxpr looks quite different:

jax.jit(counter.count).trace(0).jaxpr

{ lambda ; a:i32[]. let
    b:i32[] = add a 1:i32[]
    c:i32[] = add a 1:i32[]
  in (b, c) }

The state flows visibly through the recorded program: in as the argument, out as a result. Nothing is baked in or hidden, so every transformation handles this function correctly.

What we did to the counter works for any stateful computation. Take a class of the form

class StatefulClass:

  state: State

  def stateful_method(*args, **kwargs) -> Output:
    ...

and turn it into functions of the form

def stateless_method(state: State, *args, **kwargs) -> tuple[Output, State]:
  ...

This is a common functional programming pattern, and it’s how state is expressed in nearly all JAX programs. Training loops are the most common example: the parameters, and any optimizer state, thread through a pure update function step after step:

def update(params, opt_state, batch):
  """One training step: consume state, produce new state."""
  grads = jax.grad(loss_fn)(params, batch)
  new_params, new_opt_state = optimizer_step(params, grads, opt_state)
  return new_params, new_opt_state

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

def optimizer_step(params, grads, opt_state, lr=0.1, decay=0.9):
  # gradient descent with momentum: the momentum is optimizer state
  new_opt_state = jax.tree.map(
      lambda m, g: decay * m + g, opt_state, grads)
  new_params = jax.tree.map(
      lambda p, m: p - lr * m, params, new_opt_state)
  return new_params, new_opt_state

params = {'w': jnp.float32(1.0), 'b': jnp.float32(0.0)}
opt_state = jax.tree.map(jnp.zeros_like, params)
batch = (jnp.array([1.0, 2.0, 3.0]), jnp.array([3.0, 5.0, 7.0]))

for step in range(100):
  params, opt_state = update(params, opt_state, batch)

print(jax.tree.map(lambda x: round(float(x), 2), params))  # fits y = 2x + 1

{'b': 1.0, 'w': 2.0}

Notice that the state here is a pytree, so the pattern scales from a single counter to an entire model without changing shape (see Pytrees). This is also the convention you’ll meet everywhere in the JAX ecosystem: optimizer libraries like Optax are built around update(grads, opt_state, ...) -> (updates, new_opt_state), and neural network libraries handle parameters the same way.

Threading state as values has another benefit: because each state is an immutable snapshot, transformations apply to the whole loop. You can differentiate through an update step, or vmap it to run many independent training runs at once, without worrying about aliased mutations.

Refs: mutable arrays#

Threading values is the standard approach, but it can be awkward. If a function deep in your call stack wants to update normalization statistics or record a metric, every function along the way must plumb that state in and out of its signature.

For cases like this, JAX has refs: mutable arrays that can be read and written in place. Create one with jax.new_ref:

x_ref = jax.new_ref(jnp.zeros(3))  # new array ref, with initial value [0., 0., 0.]

@jax.jit
def f():
  x_ref[1] += 1.  # indexed add-update

print(x_ref)  # Ref([0., 0., 0.])
f()
f()
print(x_ref)  # Ref([0., 2., 0.])

Ref([0., 0., 0.], dtype=float32)
Ref([0., 2., 0.], dtype=float32)

For a ref called x_ref, we can read its entire value into an Array by writing x_ref[...], and write its entire value using x_ref[...] = A for some Array-valued expression A:

def g(x):
  x_ref = jax.new_ref(0.)
  x_ref[...] = jnp.sin(x)
  return x_ref[...]

print(jax.grad(g)(1.0))  # 0.54

0.5403023

Refs are a distinct type from Array, and come with some important constraints and limitations. In particular, indexed reading and writing is just about the only thing you can do with a ref. References can’t be passed where Arrays are expected:

x_ref = jax.new_ref(1.0)
jnp.sin(x_ref)  # error! can't do math on refs

TypeError: sin requires ndarray or scalar arguments, got <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'> at position 0.

To do math, you need to read the ref’s value first, like jnp.sin(x_ref[...]).

Reads and writes accept any NumPy indexing expression:

x_ref = jax.new_ref(jnp.arange(12.).reshape(3, 4))

# int indexing
row = x_ref[0]
x_ref[1] = row

# tuple indexing
val = x_ref[1, 2]
x_ref[2, 3] = val

# slice indexing
col = x_ref[:, 1]
x_ref[0, :3] = col

# advanced int array indexing
vals = x_ref[jnp.array([0, 0, 1]), jnp.array([1, 2, 3])]
x_ref[jnp.array([1, 2, 1]), jnp.array([0, 0, 1])] = vals

Indexing mostly follows NumPy behavior, except that an out-of-bounds index raises an error when its value is known in advance (unlike with Arrays, where reads clamp and writes drop; see Arrays and jax.numpy).

When you’re done mutating, jax.freeze invalidates the ref (so that accessing it afterwards is an error) and produces its final value as an ordinary immutable Array:

final = jax.freeze(x_ref)
final

Array([[ 1.,  1.,  9.,  3.],
       [ 1.,  3.,  2.,  3.],
       [ 9.,  9., 10.,  2.]], dtype=float32)

def normalize(x):        # pure: refs used internally only
  acc = jax.new_ref(0.0)
  acc[...] = jnp.sum(x)  # (a real program would do something less trivial)
  return x / acc[...]

jax.grad(lambda x: normalize(x).sum())(jnp.arange(1.0, 4.0))

Array([0., 0., 0.], dtype=float32)

def scale_into(x, out_ref):   # impure: takes a ref argument
  out_ref[...] = 2.0 * x

xs = jnp.arange(3.0)
out_ref = jax.new_ref(jnp.zeros(3))
jax.vmap(scale_into)(xs, out_ref)   # each instance writes its own entry
print(out_ref[...])

[0. 2. 4.]

r = jax.new_ref(0.0)

def write_shared(x):
  r[...] = x    # every batch member writes the same ref!

jax.vmap(write_shared)(jnp.arange(3.0))

Exception: performing a set/swap operation with vmapped value on an unbatched array reference of type Ref{float32[]}. Move the array reference to be an argument to the vmapped function?

Next steps#

This completes the tour of how to express computations in JAX: arrays and jax.numpy, grad and vmap, pytrees for structure, keys for randomness, and threaded values or refs for state.

Making it fast is a matter of jax.jit, sharding, and profiling, which is where the performance and scaling docs pick up: Just-in-time compilation.