triton.language.gather

triton.language.gather(src, index, axis, _semantic=None)

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

参数:

• src (Tensor) – 源张量

• index (Tensor) – 索引张量

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

该函数也可以作为 tensor 的成员函数调用，即使用 x.gather(...) 代替 gather(x, ...)。