megengine.functional.scatter#

scatter(inp, axis, index, source)[source]#

Writes all values from the tensor source into input tensor at the indices specified in the index tensor.

For each value in source, its output index is specified by its index in source for axis != dimension and by the corresponding value in index for axis = dimension.

For a 3-D tensor, input tensor is updated as:

inp[index[i][j][k]][j][k] = source[i][j][k]  # if axis == 0
inp[i][index[i][j][k]][k] = source[i][j][k]  # if axis == 1
inp[i][j][index[i][j][k]] = source[i][j][k]  # if axis == 2

inp, index and source should have same number of dimensions.

It is also required that source.shape(d) <= inp.shape(d) and index.shape(d) == source.shape(d) for all dimensions d.

Moreover, the values of index must be between 0 and inp.shape(axis) - 1 inclusive.

Note

Please notice that, due to performance issues, the result is uncertain on the GPU device if scattering different positions from source to the same destination position regard to index tensor.

Check the following examples, the oup[0][2] is maybe from source[0][2] which value is 0.2256 or source[1][2] which value is 0.5339 if set the index[1][2] from 1 to 0.

Parameters:

  • inp (Tensor) – inp tensor which to be scattered.

  • axis (int) – axis along which to index.

  • index (Tensor) – indices of elements to scatter.

  • source (Tensor) – source element(s) to scatter.

Return type:

Tensor

Returns:

output tensor.

Examples

>>> import numpy as np
>>> inp = Tensor(np.zeros(shape=(3,5),dtype=np.float32))
>>> source = Tensor([[0.9935,0.9465,0.2256,0.8926,0.4396],[0.7723,0.0718,0.5939,0.357,0.4576]])
>>> index = Tensor([[0,2,0,2,1],[2,0,1,1,2]])
>>> oup = F.scatter(inp, 0, index, source)
>>> oup.numpy()
array([[0.9935, 0.0718, 0.2256, 0.    , 0.    ],
       [0.    , 0.    , 0.5939, 0.357 , 0.4396],
       [0.7723, 0.9465, 0.    , 0.8926, 0.4576]], dtype=float32)