Skip to content

← JAX Vectorization & Control Flow step 24 of 25

Medium Primitives

lax.scan with reverse=True

Jump to the task

Accumulate from the last element

A suffix sum at position i adds x[i] and every value after it. lax.scan(body, init, xs, reverse=True) visits the last input first, then works toward the start. Each emitted value is placed at the corresponding original input index, so the output stays aligned with x.

Reverse traversal is useful whenever an earlier result depends on later values. Discounted future returns are another example, using new_carry = reward + gamma*carry instead of simple addition. Here we focus on sums, with no discount factor. The flag controls traversal; it does not request reverse-mode automatic differentiation.

Worked mini-example

from jax import lax
import jax.numpy as jnp

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

x = jnp.array([1., 2., 3.])
final, ys = lax.scan(
    step, 0.0, x,
    reverse=True
)
# final = 6.0
# ys = [6.0, 5.0, 3.0]

The steps run in this order, starting with carry zero:

Index Input New carry
2 3.0 3.0
1 2.0 5.0
0 1.0 6.0

Each new carry is stored at its original input index in ys. Reading ys from index 0 therefore gives [6, 5, 3], even though the values were computed in the order 3, 5, 6.

The body signature is identical to forward scan: (carry, x_i) → (new_carry, output). Just add reverse=True to the lax.scan call.

Keep traversal order separate from output order

Reversing x and running a forward scan gives emitted values in traversal order. For [1, 2, 3], that gives [3, 5, 6]. To reproduce reverse=True manually, you would also reverse that output, giving [6, 5, 3]. Do not reverse the output of reverse=True a second time.

The carry starts at zero before the last element is processed. It must keep its shape and dtype throughout the loop. Initialize it with jnp.zeros((), dtype=x.dtype) to match the input explicitly.

Problem

Implement reverse_cumsum(x) using lax.scan with reverse=True. The input is a one-dimensional real floating-point JAX array. Return an array of the same shape and dtype with out[i] = sum(x[i:]).

Include the current element: the last output equals the last input, and the first output is the sum of the whole array. For an empty input, return an empty array. Return all suffix sums, not the final scalar carry. Output tests check these values but cannot enforce the requested scan.