ttnn.experimental.cumsum

ttnn.experimental.cumsum = Operation(python_fully_qualified_name='ttnn.experimental.cumsum', function=<ttnn._ttnn.operations.experimental.cumsum_t object>, preprocess_golden_function_inputs=<function default_preprocess_golden_function_inputs>, golden_function=None, postprocess_golden_function_outputs=<function default_postprocess_golden_function_outputs>, is_cpp_operation=True, is_experimental=False)

Returns cumulative sum of input along dimension dim For a given input of size N, the output will also contain N elements and be such that: This function is fundamentally identical to torch.cumsum()

Parameters:

• input (ttnn.Tensor) – input tensor

• dim (int) – dimension along which to compute cumulative sum

Keyword Arguments:

• dtype (ttnn.DataType, optional) – desired output type. If specified then input tensor will be casted to dtype before processing.

• output (ttnn.Tensor, optional) – preallocated output. If specified, output must have same shape as input, and must be on the same device.

Returns:

ttnn.Tensor – the output tensor.

Note

If both dtype and output are specified then output.dtype must be dtype)

Supported dtypes, layout, ranks and dim values:

Dtypes	Layouts	Ranks	dim
BFLOAT16, FLOAT32	TILE	1, 2, 3, 4, 5	-rank <= dim < rank
INT32, UINT32	TILE	3, 4, 5	dim in {0, 1, …, rank - 3} or dim in {-rank, -rank + 1, …, -3}

Example:

import torch
import ttnn

# Create tensor
torch_input = torch.rand([2, 3, 4])
tensor_input = ttnn.from_torch(torch_input, device=device)

# Apply ttnn.experimental.cumsum() on dim=0
tensor_output = ttnn.experimental.cumsum(tensor_input, dim=0)

# With preallocated output and dtype
preallocated_output = ttnn.from_torch(torch.rand([2, 3, 4]), dtype=ttnn.bfloat16, device=device)

tensor_output = ttnn.experimental.cumsum(tensor_input, dim=0, dtype=torch.bfloat16, output=preallocated_output)