Skip to content

← JAX Vectorization & Control Flow step 5 of 25

Easy Primitives

Cumulative Sum via lax.scan

Jump to the task

Keep a total as you read the sequence

lax.scan walks through an input array and passes state from one step to the next. That state is called the carry. Each step also returns an output, which scan saves in a result array.

For a cumulative sum, the carry is the total so far. After adding the next input, return the new total twice: once as the next carry, and once as the saved output. These two roles happen to use the same value here; other scans can carry one value and emit another.

With its default unroll setting, scan uses a compiled loop rather than expanding a Python loop inside JIT tracing. This can reduce compilation cost for long sequences; it does not guarantee faster execution for every workload.

Worked mini-example

import jax.numpy as jnp
from jax import lax

def step(carry, x_i):
    new = carry + x_i
    return new, new

x = jnp.array([1., 2., 3., 4.])
final, ys = lax.scan(
    step, 0.0, x
)
# final = 10.0
# ys = [1., 3., 6., 10.]

lax.scan(body, init, xs) returns (final_carry, stacked_outputs). Here the final carry is a scalar, 10, and the outputs have shape (4,). The same running total serves two roles at each step: it is passed to the next step and saved in the output array.

Step through the sum

Choose Step to consume one input or Run all to finish. The addition mode starts with carry zero and uses the same four inputs as the code above. Other operations let you explore how the body changes the result; Reset starts the selected operation again.

Loading visualization…

Common pitfalls

  • Returning a single value from body. The body must return a pair (new_carry, output). Returning just new_carry is not a valid body result.
  • Changing the carry's dtype. The carry must keep its shape and dtype. An explicitly typed integer carry can conflict with floating inputs. A Python scalar can be promoted, so it is not accurate to say that plain 0 always fails. jnp.zeros((), dtype=x.dtype) makes the intended scalar dtype explicit.
  • Confusing the return order. lax.scan returns (final_carry, outputs). drop the carry with _ if you only need the per-step outputs.

Problem

Implement cumulative_sum(x) using lax.scan. Do not use jnp.cumsum. the point is to learn the scan API.

The input is a one-dimensional real floating-point JAX array. Return an array of the same shape and dtype, containing each prefix sum: out[i] = sum(x[:i+1]). Include the current element, so the first output is x[0], not zero. An empty input returns an empty array.

Return the stacked outputs, not just the final total. Output tests verify values but cannot enforce scan over an equivalent implementation.