Skip to content

← JAX Vectorization & Control Flow step 23 of 25

Hard Primitives

Associative Scan (Parallel Cumsum)

Jump to the task

Group the work without changing the intended operation

A prefix sum returns the total up to each input position. A sequential scan extends one total at a time. An associative scan combines groups of values, allowing independent combinations to run in parallel.

This regrouping is valid when the operation is associative: op(op(a, b), c) = op(a, op(b, c)). For exact addition, (1 + 2) + 3 and 1 + (2 + 3) both give 6. Associativity concerns grouping, not reordering. An operation can be associative without being commutative, as matrix multiplication illustrates.

For a constant-cost combine operation, a parallel prefix algorithm can have logarithmic dependency depth instead of a sequential chain of N additions. That does not mean logarithmic total work or a guaranteed wall-clock speedup. Sequence length, hardware and overhead matter. A regular scan can also use parallel work inside each step.

Compute every prefix

from jax import lax
import jax.numpy as jnp

x = jnp.array([1.0, 2.0, 3.0, 4.0])
out = lax.associative_scan(lambda a, b: a + b, x)
# [1.0, 3.0, 6.0, 10.0]

Unlike lax.scan, this API has no explicit initial carry and does not return a (carry, outputs) pair. It returns the prefix results directly. The first result is x[0]; each later result includes its corresponding input. The default scan axis is 0.

The combine function receives arrays of partial results and must combine corresponding entries, preserving their shape. Addition already does this. Reducing a + b to a scalar would break that contract. For matrix-valued elements, batched matrix multiplication can combine corresponding matrices without changing their order.

Floating-point addition is not exactly associative

Finite precision can change the answer when additions are regrouped. In float32, (1e8 + -1e8) + 1 gives 1, while 1e8 + (-1e8 + 1) gives 0: the inner addition loses the small value. Parallel and sequential prefix sums can therefore differ, sometimes substantially relative to the final result. Do not assume bitwise agreement or that cancellation is always harmless.

Subtraction is not associative even in exact arithmetic: (5 - 2) - 1 = 2, but 5 - (2 - 1) = 4. JAX does not prove your combine function is associative. Choosing subtraction can silently produce a different result from the sequential operation you intended.

Problem

Implement parallel_cumsum(x) with lax.associative_scan and elementwise addition. Do not use jnp.cumsum or lax.scan for this exercise. The input is a one-dimensional real floating-point JAX array. Return an array with the same shape and dtype, containing each prefix sum: out[i] = sum(x[:i+1]), subject to floating-point rounding. An empty input returns an empty array.

Return all prefix sums, not just the total. Output tests check the values but cannot enforce the requested API over an equivalent implementation.