Skip to content

← JAX Vectorization & Control Flow step 6 of 25

Medium Primitives

Running Mean via lax.scan

Jump to the task

Carry a sum and a count

A prefix mean averages all values seen so far. For [4, 2, 6], the successive means are [4, 3, 4]. This is not a fixed-window moving average or an exponentially weighted average: each prefix includes every earlier value with equal weight.

To compute it here, carry a pair (sum, count) through the sequence. Add the new value to the sum and increment the count, then emit their ratio. JAX accepts tuples and other pytrees as scan carries, provided the structure and each leaf's shape and dtype stay fixed across steps.

Worked mini-example

from jax import lax
import jax.numpy as jnp

def step(carry, x_i):
    s, n = carry
    s = s + x_i
    n = n + 1
    return (s, n), s/n

x = jnp.array([4., 2.])
final, means = lax.scan(
    step, (0.0, 0), x
)
# final = (6.0, 2)
# means = [4.0, 3.0]

The carry starts at (0.0, 0): a scalar sum and a scalar count. Both are updated before computing the mean:

Input Carry Mean
4.0 (4.0, 1) 4.0
2.0 (6.0, 2) 3.0

The pair is passed to the next step, while the scalar mean is saved in the output array. means has the same shape as x; the carry remains a pair of scalars rather than growing with the sequence.

Common pitfalls

  • Changing the carry. Returning a different structure, leaf shape or dtype from one step to the next is an error. A consistently floating-point count can work too; the requirement is consistency, not that the count must always have an integer dtype.
  • Forgetting the count update. If n remains at its initial zero, division produces non-finite values, not a running mean. Update both sum and count before dividing.
  • Dividing too soon. Dividing by the old count makes the first step divide by zero. Python 0 is a valid initial count here; an explicit jnp.int32(0) is also valid but is not required for JIT.

Problem

Implement running_mean(x) using lax.scan with a tuple carry (running_sum, count). Do not use jnp.cumsum or Python loops.

The input is a one-dimensional real floating-point JAX array. Return an array of the same shape containing out[i] = mean(x[:i+1]). The current value belongs in both the sum and the count. An empty input returns an empty array; no mean needs to be evaluated.

Start the sum with jnp.zeros((), dtype=x.dtype) and the count with zero. Floating-point summation can round or overflow, especially for long sequences or small dtypes. This exercise teaches the scan state, not a compensated summation algorithm.

Output tests check the prefix means, but cannot enforce the requested tuple carry or distinguish an equivalent implementation using cumsum.