sponge.function.gather_vector

查看源文件
sponge.function.gather_vector(tensor: Tensor, index: Tensor)[源代码]

根据指标从张量的倒数第二轴收集向量。

参数:
  • tensor (Tensor) - 输入张量,shape为 \((B, X, D)\) 。其中 \(B\) 为batch size, \(X\) 为任意大小, \(D\) 为模拟系统的维度,通常为3。

  • index (Tensor) - 索引,shape为 \((B, ...,)\)

返回:

Tensor。取出的向量。shape为 \((B, ..., D)\)

支持平台:

Ascend GPU