oneflow.scatter_add

oneflow.scatter_add(input, dim, index, src)

This operator scatter the src with addition operation according to index along dim into the input.

Take a 3-D blob as example, the output is specified by:

input[index[i][j][k]][j][k] += src[i][j][k]  # if dim == 0
input[i][index[i][j][k]][k] += src[i][j][k]  # if dim == 1
input[i][j][index[i][j][k]] += src[i][j][k]  # if dim == 2
Parameters
  • input (Tensor) – The input blob.

  • dim (int) – The axis along which to index

  • index (Tensor) – The index blob of elements to scatter.

  • src (Tensor) – The source blob whose elements will be scatterd and added to output.

Returns

The scatterd Tensor.

Return type

Tensor

For example:

>>> import oneflow as flow
>>> import numpy as np
>>> input = flow.ones((3,5))*2
>>> index = flow.tensor(np.array([[0,1,2],[0,1,4]], ), dtype=flow.int32)
>>> src = flow.Tensor(np.array([[0,10,20,30,40],[50,60,70,80,90]]))
>>> out = flow.scatter_add(input, 1, index, src)
>>> out
tensor([[ 2., 12., 22.,  2.,  2.],
        [52., 62.,  2.,  2., 72.],
        [ 2.,  2.,  2.,  2.,  2.]], dtype=oneflow.float32)