triton.language.cumsum

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

返回沿给定 axis 的 input 张量中所有元素的累积和 (cumsum)

参数:

• input (Tensor) – 输入值

• axis (int) – 执行扫描操作的维度

此函数也可以作为成员函数在 tensor 上调用，使用 x.cumsum(...) 而不是 cumsum(x, ...) 的方式。