Skip to content

← JAX Foundations step 1 of 25

Easy End-to-End

Pure Function vs Impure

Jump to the task

Why this matters

A pure function gets the values it needs from its inputs and returns its results without changing hidden external state. To update a running history, pass the current history in and return its replacement.

This matters for JIT because Python side effects happen while JAX traces the function. They are not repeated each time the compiled computation runs. Appending an input to a global list can even leave a tracer in that list instead of a concrete value.

Worked mini-example

import jax.numpy as jnp

def append_and_average(state, x):
    new_state = jnp.concatenate([state, x])
    return new_state, jnp.mean(new_state)

state = jnp.array([1.0, 3.0])
state, mean = append_and_average(state, jnp.array([5.0]))
# state = [1, 3, 5], mean = 3

state, mean = append_and_average(state, jnp.array([7.0]))
# state = [1, 3, 5, 7], mean = 4

Each call returns a new array. The caller keeps it for the next call. Earlier arrays are unchanged.

Growing a history changes its shape, so a jitted version may need another compilation for each new length. This is a simple purity example, not an efficient design for an indefinitely growing history.

Common pitfalls

  • Pass changing values as inputs. A global or captured value is not automatically refreshed when a compiled function is reused.
  • Local lists and dictionaries are fine when they stay within the call. Using a class is also fine if changing state is explicit.
  • Ordinary Python print runs during tracing. Use jax.debug.print when you need values from compiled execution.
  • The exercise packs the updated state and mean together. Before the next call, take packed[:-1] as the state; the last element is the mean and must not become another history item.

Problem

Implement make_pure(state, x). Both inputs are one-dimensional floating-point JAX arrays. x contains exactly one element; state may be empty.

Append x to state, calculate the mean of that updated history, then append the mean to the output. Return one array of shape (state.shape[0] + 2,), leaving both inputs unchanged.

For state = [1, 3] and x = [5], return [1, 3, 5, 3]. For an empty state and x = [5], return [5, 5].

The single-array format is this exercise's contract. A tuple (new_state, mean) is also valid in JAX and is often clearer in an application.