Skip to main content

Mojo function

cumsum

cumsum[dtype: DType, exclusive: Bool, reverse: Bool](output: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], input: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], axis: Int)

Implements the CumSum operator from the ONNX spec: https://github.com/onnx/onnx/blob/main/docs/Operators.md#CumSum Computes cumulative sum of the input elements along the given axis. Cumulative sum can be inclusive or exclusive of the top element, and normal or reverse (direction along a given axis).

Parameters:

  • dtype (DType): Type of the input and output tensors.
  • exclusive (Bool): If set to True, return exclusive sum (top element not included).
  • reverse (Bool): If set to True, perform cumsum operation in reverse direction.

Args:

  • output (LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]): The output tensor.
  • input (LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]): The input tensor.
  • axis (Int): The axis on which to perform the cumsum operation.

Was this page helpful?