PyTorch torch.cumsum Function
Pytorch torch Reference Manual
torch.cumsumis a function in PyTorch used to compute the cumulative sum. It returns the cumulative sum along a specified dimension, i.e., the sum of all elements from the beginning to the current position.
Function Definition
torch.cumsum(input, dim, dtype=None)
Usage Example
Example
import torch
# Compute the cumulative sum
x = torch.tensor([1, 2, 3, 4, 5])
result = torch.cumsum(x, dim=0)
print("Input:", x)
print("Cumulative sum:", result)
# Output: tensor([1, 3, 6, 10, 15])
# Explanation: 1, 1+2=3, 1+2+3=6, 1+2+3+4=10, 1+2+3+4+5=15
# 2-dimensional tensor
x = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)
# Cumulative along columns
result_col = torch.cumsum(x, dim=0)
print("\nCumulative sum along columns:")
print(result_col)
# tensor([[ 1, 2, 3],
# [ 5, 7, 9]])
# Cumulative along rows
result_row = torch.cumsum(x, dim=1)
print("\nCumulative sum along rows:")
print(result_row)
# tensor([[ 1, 3, 6],
# [ 4, 9, 15]])
# Compute the cumulative sum
x = torch.tensor([1, 2, 3, 4, 5])
result = torch.cumsum(x, dim=0)
print("Input:", x)
print("Cumulative sum:", result)
# Output: tensor([1, 3, 6, 10, 15])
# Explanation: 1, 1+2=3, 1+2+3=6, 1+2+3+4=10, 1+2+3+4+5=15
# 2-dimensional tensor
x = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)
# Cumulative along columns
result_col = torch.cumsum(x, dim=0)
print("\nCumulative sum along columns:")
print(result_col)
# tensor([[ 1, 2, 3],
# [ 5, 7, 9]])
# Cumulative along rows
result_row = torch.cumsum(x, dim=1)
print("\nCumulative sum along rows:")
print(result_row)
# tensor([[ 1, 3, 6],
# [ 4, 9, 15]])
Other Extensions