Running totals

A cumulative sum records the total up to each position rather than only the final total. It preserves the input shape.

import torch

x = torch.tensor([[1.0, 2.0, 3.0, 4.0], [10.0, -2.0, 1.0, 0.0]])
print(x.cumsum(dim=-1).tolist())
# [[1.0, 3.0, 6.0, 10.0], [10.0, 8.0, 9.0, 9.0]]

Each row starts its own running total. Mathematically, the last entry equals the row sum. Floating-point implementations can round differently if they combine values in a different order.

Use the tensor operation

A Python loop can calculate the same running totals, but indexing, extracting Python scalars and writing individual elements add overhead. On a GPU, scalar extraction can also wait for device work.

Prefix sums admit parallel algorithms, so a sequential-looking definition does not require a Python loop. cumsum uses an implementation selected for the device and input. Do not assume every backend uses one particular algorithm or that fewer dispatched operations guarantee a given speedup.

Related operations

diff computes adjacent differences. To reconstruct values from those differences, you also need the initial value. cumprod, cummax and cummin accumulate products, maxima and minima respectively. logcumsumexp computes the logarithm of cumulative sums of exponentials; it is not a replacement for cumulative products.

Cumulative sums are useful for running counts and converting sequence lengths into offsets. A reversed cumulative sum gives undiscounted future reward totals. Discounted returns require discount factors as well.

Give each tensor axis a meaning

A tensor's shape tells you how many entries each axis has. It does not tell you what those entries represent. For example, (batch, features) and (features, batch) use different axis orders. Name the axes before choosing an operation.

A reduction combines entries along an axis, for example by summing or averaging. Averaging the features of a (B, D) tensor gives one number per example, with shape (B,). Keeping the reduced axis gives shape (B, 1) instead.

For [[2, 4, 6], [1, 3, 5]], the row means with the axis kept are [[4], [3]]. Subtracting them from the input gives [[-2, 0, 2], [-2, 0, 2]]: each row has been centered using its own mean.

Reshape and transpose do different jobs

A reshape regroups elements using the operation's traversal order. A transpose changes the axis order. For the usual row-major reshape, [[1, 2, 3], [4, 5, 6]] reshaped to (3, 2) becomes [[1, 2], [3, 4], [5, 6]]. Transposing the original instead gives [[1, 4], [2, 5], [3, 6]].

Both results have shape (3, 2), but their entries mean different things. Check values as well as shapes. Unequal axis sizes make many axis mistakes easier to spot than square examples do.

Concatenate along an axis or stack on a new one

Concatenation extends an existing axis. Stacking creates a new axis. For two arrays of shape (2, 3), concatenating on axis zero gives (4, 3). Stacking on a new leading axis gives (2, 2, 3).