快捷方式

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

inputindex 必须具有相同数量的维度。还需要 index.size(d) <= input.size(d) 对于所有维度 d != dimout 将与 index 形状相同。请注意,inputindex 不会互相广播。

参数
  • input (Tensor) – 源张量

  • dim (int) – 要索引的轴

  • index (LongTensor) – 要收集的元素的索引

关键字参数
  • sparse_grad (bool, 可选) – 如果为 True,则针对 input 的梯度将是一个稀疏张量。

  • out (Tensor, 可选) – 目标张量

示例

>>> t = torch.tensor([[1, 2], [3, 4]])
>>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]]))
tensor([[ 1,  1],
        [ 4,  3]])

文档

访问 PyTorch 的综合开发者文档

查看文档

教程

获取针对初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源