triton.language.gather

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

沿给定维度从张量中收集数据。

Parameters:

• src (Tensor) – 源张量

• index (Tensor) – 索引张量

• axis (int) – 要沿其聚集的维度

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