triton.language.associative_scan(input, axis, combine_fn)

Applies the combine_fn to each elements with a carry in input tensors along the provided axis and update the carry

  • input – the input tensor, or tuple of tensors

  • axis – the dimension along which the reduction should be done

  • combine_fn – a function to combine two groups of scalar tensors (must be marked with @triton.jit)