triton.language.cumsum

triton.language.cumsum(input, axis=0, reverse=False)

Returns the cumsum of all elements in the input tensor along the provided axis

Parameters:
  • input (Tensor) – the input values

  • axis (int) – the dimension along which the scan should be done

This function can also be called as a member function on tensor, as x.cumsum(...) instead of cumsum(x, ...).