triton.language.cumsum

triton.language.cumsum(input, axis=0, reverse=False, dtype: constexpr = None)

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

参数:

• input (Tensor) – 输入值

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

• reverse (bool) – 如果为 true，则沿相反方向执行扫描

• dtype (tl.dtype) – 返回张量所需的预期数据类型。如果指定了该参数，输入张量将在执行操作前被转换为 dtype。如果未指定，较小的整数类型（< 32 位）会被自动提升以防止溢出。请注意，tl.bfloat16 类型的输入会自动提升为 tl.float32。

该函数也可以作为 tensor 的成员函数调用，即使用 x.cumsum(...)，而不是 cumsum(x, ...)。