Running totals

A reduction collapses an axis. A scan keeps it, and stores every partial result along the way.

torch.tensor([1., 2., 3., 4.]).cumsum(0)     # [1., 3., 6., 10.]

The last entry is what sum() would have given. The rest is the working.

It looks sequential and is not

Each entry depends on the one before it, so the obvious implementation is a loop and the obvious conclusion is that it cannot be parallelised. Both are wrong: a prefix sum is computed in log(n) parallel steps by a standard algorithm, which is what cumsum dispatches. Writing the loop yourself throws that away and adds a dispatch per element.

This is worth internalising beyond this function. “Each step depends on the last” is not by itself a reason to write a Python loop.

The inverse

torch.diff(x)      # [x1 - x0, x2 - x1, ...], one shorter

diff and cumsum undo each other up to the first element, which makes the pair useful for converting between rates and totals: gradients to positions, deltas to prices, per-step rewards to returns.

The family

cumsum, cumprod, cummax, cummin, and logcumsumexp for when the products would overflow, on exactly the reasoning from the softmax problem.

Where you will use it

Sequence lengths to offsets when packing ragged batches. Discounted returns in reinforcement learning, as a cumsum on reversed rewards. Any “so far” quantity over a time axis.

Your task

def running_total(x: torch.Tensor) -> torch.Tensor

Return the running total along the last axis, in at most two dispatched operations.