torch.gather¶
- torch.gather(input, dim, index, *, sparse_grad=False, out=None) Tensor ¶
沿由 dim 指定的轴收集值。
对于一个 3 维张量,输出由
out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0 out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1 out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2
input
和index
必须具有相同数量的维度。还需要index.size(d) <= input.size(d)
对于所有维度d != dim
。out
将与index
形状相同。请注意,input
和index
不会互相广播。- 参数
- 关键字参数
示例
>>> t = torch.tensor([[1, 2], [3, 4]]) >>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]])) tensor([[ 1, 1], [ 4, 3]])