torch.cumsum(): Cumulative Sum of Elements in a Tensor

The torch.cumsum() method in PyTorch calculates the cumulative sum of elements along a specified dimension of a tensor. It returns a tensor where each element is the sum of all elements up to and including the current position in the specified dimension.

Syntax

torch.cumsum(input, dim, dtype=None, out=None)

Parameters

Argument	Description
input (Tensor)	It represents an input tensor.
dim (int)	It is a dimension to calculate the cumulative sum.
dtype (torch.dtype, optional)	It is the desired data type of the output. Default: same as input.
out (Tensor, optional)	It is an output tensor if defined. By default, it is None.

Basic 1D Tensor

Let’s find out the cumulative sum of elements of the input 1D tensor.

If input is 1D, dim=0 calculates the cumulative sum across the entire tensor.

import torch

tensor = torch.tensor([1, 2, 3, 4])

cum_sum = torch.cumsum(tensor, dim=0)

print(cum_sum)

# Output: tensor([ 1,  3,  6, 10])
# tensor([0+1, 0+1+2, 0+1+2+3, 0+1+2+3+4])

Here is how we reached the output tensor:

1. Step 0: 1
2. Step 1: 1 + 2 = 3
3. Step 2: 1 + 2 + 3 = 6
4. Step 3: 1 + 2 + 3 + 4 = 10

You can see that the output tensor has the same shape as the input.

2D Tensor (Along rows vs. columns)

Cumulative sum along rows

For calculating the cumulative sum along rows, use the dim=1 argument.

import torch

tensor_2d = torch.tensor([[5, 6, 7],
                          [1, 2, 3]])

cum_sum_row = torch.cumsum(tensor_2d, dim=1)

print(cum_sum_row)

# Output: tensor([[ 5, 11, 18],
#                 [ 1,  3,  6]])

In this code, we are performing an accumulating sum row-wise. 

For the first row, 5, 0+ 5 + 6 = 11, 5 + 6 + 7 = 18: [5, 11, 18].

For the second row, 1, 0+1+2 = 3, 0+1+2+3 = 6: [1, 3, 6].

Cumulative sum along columns

For calculating the cumulative sum along columns, use the dim=0 argument.

import torch

tensor_2d = torch.tensor([[5, 6, 7],
                          [1, 2, 3]])

cum_sum_columns = torch.cumsum(tensor_2d, dim=0)

print(cum_sum_columns)

# Output: 
# tensor([[ 5,  6,  7],
#         [ 6,  8, 10]])

For each column, the cumulative sum is calculated: Column 1: [5, 1+5=6], Column 2: [6, 6+2=8], Column 3: [7, 7+3=10].

Specifying Data Type

For controlling the data type of the output tensor, use the “dtype” argument.

import torch

tensor = torch.tensor([1.5, 2.5, 3.5])

cumulative_int = torch.cumsum(tensor, dim=0, dtype=torch.int32)

print(cumulative_int)

# Output: tensor([1, 3, 6], dtype=torch.int32)

Using the output tensor

If we have a pre-allocated tensor, we can store the result in it using the “out” argument.

import torch

x = torch.tensor([1, 2, 3])

out = torch.zeros(3, dtype=torch.int64)

torch.cumsum(x, dim=0, out=out)

print(out)

# Output: tensor([1, 3, 6])

It avoids creating a new tensor, which can sometimes be helpful in terms of performance.