torch.gather/torch.scatter

torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor
参数：

• input 被索引的tensor
• dim 索引沿着的维度
• index 索引的index

官方给出的三维例子：

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

可以看出以下几点：

• 输入的index的每一个值都会被作为输入dim的index来取值
• index的ndims和input的ndims相同
• 结果的shape与index的shape相同

写一个二维版理解下：

out[i][j] = input[index[i][j]][j]  # if dim == 0
out[i][j] = input[i][index[i][j]]  # if dim == 1

可以看出隐藏的限制：

• dim=0，index的维度1应小于等于input的维度1
• dim=1， index的维度0应小于等于input的维度0

>>>p=torch.randn(4,2)
>>>p
tensor([[ 0.7786,  1.4472],
        [ 0.6529, -0.0105],
        [ 0.8745,  0.0016],
        [-0.8376, -0.4966]])

>>>p.gather(0,torch.tensor([[1,0]]))
tensor([[0.6529, 1.4472]])

>>>p.gather(0,torch.tensor([[1,0,1]]))
RuntimeError: Size does not match at dimension 1 expected index [1, 3] to be smaller than src [4, 2] apart from dimension 0

>>>p.gather(0,torch.tensor([[1,0],[0,1],[1,1],[0,0],[1,1]]))
tensor([[ 0.6529,  1.4472],
        [ 0.7786, -0.0105],
        [ 0.6529, -0.0105],
        [ 0.7786,  1.4472],
        [ 0.6529, -0.0105]])

>>>p.gather(1,torch.tensor([[1,0],[0,1]]))
tensor([[ 1.4472,  0.7786],
        [ 0.6529, -0.0105]])
  
>>>p.gather(1,torch.tensor([[1,0,1],[0,1,1]]))
tensor([[ 1.4472,  0.7786,  1.4472],
        [ 0.6529, -0.0105, -0.0105]])

>>>p.gather(1,torch.tensor([[1],[0],[1],[1],[1]]))
RuntimeError: Size does not match at dimension 0 expected index [5, 1] to be smaller than src [4, 2] apart from dimension 1

torch官方文档 TORCH.GATHER

torch.gather是取，torch.scatter则是放
Tensor.scatter(dim, index, src, reduce=None) → Tensor

以index的值作为tensor的dim维度的索引，从src里取值

self[index[i][j][k]][j][k] = src[i][j][k]  # if dim == 0
self[i][index[i][j][k]][k] = src[i][j][k]  # if dim == 1
self[i][j][index[i][j][k]] = src[i][j][k]  # if dim == 2

可以看出有以下限制：

• tensor,index,src的ndims相同
• index.size(d) <= src.size(d)
• 当d!=dim index.size(d) <= tensor.size(d)
• 如果src为数值，则默认维度与tensor一致

>>>x = torch.zeros(5,3)
>>>index = torch.tensor([0,1,2,0,1])
>>>x.scatter_(1,index.view(-1,1),1)
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.],
        [1., 0., 0.],
        [0., 1., 0.]])

>>>index2 = torch.tensor([[0,1],[1,2],[2,0],[0,1],[1,2]])
>>>x.scatter_(1,index2,1)
tensor([[1., 1., 0.],
        [0., 1., 1.],
        [1., 0., 1.],
        [1., 1., 0.],
        [0., 1., 1.]])
        
>>>index3 = torch.randint(1,3,(5,4))
>>>x.scatter_(1,index3,1)
Traceback (most recent call last):
  File "<input>", line 1, in <module>
RuntimeError: Expected index [5, 4] to be smaller than self [5, 3] apart from dimension 1 and to be smaller size than src [5, 3]
>>>src = torch.ones(5,4)
>>>x.scatter_(1,index3,src)
tensor([[1., 1., 1.],
        [0., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.],
        [0., 1., 1.]])

TORCH.TENSOR.SCATTER_