cumsum¶

tripy.cumsum(input: Tensor, dim: int) → Tensor[source]¶

Computes the cumulative sum of elements in the input along the dimension dim.

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
dim (int) – The dimension along which to compute the cumulative sum.

Returns:

[dtype=T1] A tensor of the same shape as the input.

Return type:

Tensor

TYPE CONSTRAINTS:

T1: bfloat16, float16, float32, int32

Example: 1D tensor

1D tensor¶

input = tp.arange(4, 0, step=-1, dtype=tp.int32)
output = tp.cumsum(input, dim=0)

>>> input
tensor([4, 3, 2, 1], dtype=int32, loc=gpu:0, shape=(4,))
>>> output
tensor([4, 7, 9, 10], dtype=int32, loc=gpu:0, shape=(4,))

Example: 2D tensor

2D tensor¶

input = tp.reshape(tp.arange(9, 0, step=-1, dtype=tp.int32), (3, 3))
output = tp.cumsum(input, dim=0)

>>> input
tensor(
    [[9, 8, 7],
     [6, 5, 4],
     [3, 2, 1]], 
    dtype=int32, loc=gpu:0, shape=(3, 3))
>>> output
tensor(
    [[9, 8, 7],
     [15, 13, 11],
     [18, 15, 12]], 
    dtype=int32, loc=gpu:0, shape=(3, 3))