triton.language.gather

triton.language.gather(src, index, axis)

Gather from a tensor along a given dimension.

Parameters:

• src (Tensor) – the source tensor

• index (Tensor) – the index tensor

• axis (int) – the dimension to gather along

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