triton.language.associative_scan

triton.language.associative_scan(input, axis, combine_fn, reverse=False, _semantic=None, _generator=None)

沿提供的 axis（轴）将 combine_fn 应用于每个带有进位（carry）的 input 张量元素，并更新进位。

参数:

• input (Tensor) – 输入张量，或张量元组

• axis (int) – 执行归约操作的维度

• combine_fn (Callable) – 用于组合两组标量张量的函数（必须使用 @triton.jit 标记）

• reverse (bool) – 是否沿轴反向执行结合扫描（associative scan）

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