torch.scatter_

torch.scatter_

Tensor.scatter_(dim, index, src, reduce=None) → Tensor

Parameters

• dim (int) – the axis along which to index
• index (LongTensor) – the indices of elements to scatter, can be either empty or of the same dimensionality as src. When empty, the operation returns self unchanged.
• src (Tensor or float) – the source element(s) to scatter. 要填进去的元素
• reduce (str, optional) – reduction operation to apply, can be either 'add' or 'multiply'. 用的相对较少。

直接看例子，

>>> src = torch.arange(1, 11).reshape((2, 5))
>>> src
tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])
>>> index = torch.tensor([[0, 1, 2, 0]])
>>> torch.zeros(3, 5, dtype=src.dtype).scatter_(0, index, src)
tensor([[1, 0, 0, 4, 0],
        [0, 2, 0, 0, 0],
        [0, 0, 3, 0, 0]])
# 从这个例子出发来简单说明：首先dim=0,意味着需要沿着axis=0的方向进行操作，即index每一列逐渐增大，按列找到对应的索引号，然后按顺序把src中的元素填进去。

>>> index = torch.tensor([[0, 1, 2], [0, 1, 4]])
>>> torch.zeros(3, 5, dtype=src.dtype).scatter_(1, index, src)
tensor([[1, 2, 3, 0, 0],
        [6, 7, 0, 0, 8],
        [0, 0, 0, 0, 0]])
# dim=1, 按行找到对应的index，按顺序把src中的元素填进去

>>> torch.full((2, 4), 2.).scatter_(1, torch.tensor([[2], [3]]),1.23)
tensor([[2.0000, 2.0000, 1.2300, 2.0000],
        [2.0000, 2.0000, 2.0000, 1.2300]])
# dim=1, 按行找到对应的index，按顺序把src中的元素填进去,不用管原来的位置是什么数字。

**注意：**index可以不用满，src按顺序填充。

>>> torch.full((2, 4), 2.).scatter_(1, torch.tensor([[2], [3]]),
...            1.23, reduce='multiply')
tensor([[2.0000, 2.0000, 2.0000*1.23, 2.0000],
        [2.0000, 2.0000, 2.0000, 2.000*1.23]])
tensor([[2.0000, 2.0000, 2.4600, 2.0000],
        [2.0000, 2.0000, 2.0000, 2.4600]])
# dim=1, 按行找到对应的index，按顺序把src中的元素乘上去

>>> torch.full((2, 4), 2.).scatter_(1, torch.tensor([[2], [3]]),
...            1.23, reduce='add')
tensor([[2.0000, 2.0000, 2.0000+1.23, 2.0000],
        [2.0000, 2.0000, 2.0000, 2.000+1.23]])
tensor([[2.0000, 2.0000, 3.2300, 2.0000],
        [2.0000, 2.0000, 2.0000, 3.2300]])
# dim=1, 按行找到对应的index，按顺序把src中的元素加上去上去

Reference:

[1] TORCH.TENSOR.SCATTER_

tps://pytorch.org/docs/stable/generated/torch.Tensor.scatter_.html)