前言

我觉得pytorch最难的两个api之一就是gather和scatter了，但是在项目中又经常出现，OMG…
还是想办法弄懂吧

torch.gather

input = torch.randn((2, 3, 4))

tensor([[[ 0.1815, -0.7603,  0.3465,  0.2593],
         [ 0.5841,  1.2581,  0.1583, -0.7283],
         [-0.0093,  0.6214, -0.3391, -2.4973]],

        [[-0.8153, -0.3802, -1.9930, -2.1965],
         [-0.8470,  0.5285,  0.2684, -0.2982],
         [-0.6088, -0.6239,  1.1218, -0.1660]]])

index = torch.tensor([[[0, 0, 1]]]).long()
torch.gather(input, 0, index)

tensor([[[ 0.1815, -0.7603, -1.9930]]])

index = torch.tensor([[[0, 2, 1]]]).long()
torch.gather(input, 1, index)

tensor([[[0.1815, 0.6214, 0.1583]]])

index = torch.tensor([[[0, 2, 1]]]).long()
torch.gather(input, 2, index)

tensor([[[ 0.1815,  0.3465, -0.7603]]])

torch.scatter