Understand torch.scatter_()

1. Official Documentation

First, note that scatter_() is an inplace function, meaning that it will change the value of input tensor.

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

2. Graphical Diagram for dim=0

For simplicity, let us consider two-dimensional matrices here. Let us first understand dim.

import torch
import numpy as np
src = torch.from_numpy(np.arange(1, 11)).float().view(2, 5)
print(src)
> tensor([[ 1., 2., 3., 4., 5.],
[ 6., 7., 8., 9., 10.]])
input_tensor = torch.zeros(3, 5)
print(input_tensor)
> tensor([[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.]])
index_tensor = torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]])
print(index_tensor)
> tensor([[0, 1, 2, 0, 0],
[2, 0, 0, 1, 2]])
## try to manually work out the result
dim = 0
print(input_tensor.scatter_(dim, index_tensor, src))
> ...
> tensor([[ 1.,  7.,  8.,  4.,  5.],
[ 0., 2., 0., 9., 0.],
[ 6., 0., 3., 0., 10.]])
input_tensor.shape[dim] > index_tensor.max().item()

3. Graphical Diagram for dim = 1

Similarly, we can work it out when dim=1. Let us try the following example.

src = torch.from_numpy(np.arange(1, 11)).float().view(2, 5)
input_tensor = torch.zeros(3, 5)
index_tensor = torch.tensor([[3, 0, 2, 1, 4], [2, 0, 1, 3, 1]])
dim = 1
print(input_tensor.scatter_(dim, index_tensor, src))
> tensor([[ 2.,  4.,  3.,  1.,  5.],
[ 7., 10., 6., 9., 0.],
[ 0., 0., 0., 0., 0.]])

4. Graphical Diagram for a Trickier Example

Finally, let’s try a trickier example where the src is a value and the size of the index tensor is smaller than the input tensor for dim != dim.

input_tensor = torch.from_numpy(np.arange(1, 16)).float().view(3, 5) # dim is 2# unsqueeze to have dim = 2
index_tensor = torch.tensor([4, 0, 1]).unsqueeze(1)
src = 0
dim = 1
print(input_tensor.scatter_(dim, index_tensor, src))
> tensor([[ 1.,  2.,  3.,  4.,  0.],
[ 0., 7., 8., 9., 10.],
[11., 0., 13., 14., 15.]])

A Ph.D. student in Statistics and NLP.

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store