megengine.functional.gather¶
-
gather(inp, axis, index)[源代码]¶ 根据给定的索引从输入 Tensor 中收集数据。
对于三维 Tensor, 输出由下面的内容确定：
out[i][j][k] = inp[index[i][j][k]][j][k] # if axis == 0 out[i][j][k] = inp[i][index[i][j][k]][k] # if axis == 1 out[i][j][k] = inp[i][j][index[i][j][k]] # if axis == 2
如果
inp是一个尺寸为 \((x_0,x_1,...,x_{i-1},x_i,x_{i+1},...,x_{n-1})\) 且 axis=i 的 n 维 Tensor 则index必须是一个尺寸为 \((x_0,x_1,...,x_{i-1},y,x_{i+1},...,x_{n-1})\) 的 n 维 Tensor，这里的 \(y\ge 1\) 和输出的尺寸都必须必须与index的尺寸相同。- 参数
inp (
Tensor) – 输入张量。axis (
int) – 将要进行索引的轴。index (
Tensor) – 将要进行聚合的元素的索引
- 返回类型
Tensor- 返回
输出张量。
例如：
import megengine.functional as F from megengine import tensor inp = tensor([ [1,2], [3,4], [5,6], ]) index = tensor([[0,2], [1,0]]) oup = F.gather(inp, 0, index) print(oup.numpy())
输出：
[[1 6] [3 2]]