triton.language.gather

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

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

参数:

• src (Tensor) – 源张量

• index (Tensor) – 索引张量

• axis (int) – 要收集数据的维度

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