oneflow.scatter

oneflow.scatter(input, dim, index, src, *, reduce=None)

This operator writes the elements specified by index along with the axis dim from the src 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

input, index and src (if it is a Tensor) should all have the same number of dimensions. It is also required that index.shape(d) <= src.shape(d) for all dimensions d, and that index.shape(d) <= input.shape(d) for all dimensions d != dim. Note that index and src do not broadcast.

Warning

When indices are not unique, the behavior is non-deterministic (one of the values from src will be picked arbitrarily) and the gradient will be incorrect (it will be propagated to all locations in the source that correspond to the same index)!

Note

The backward pass is implemented only for src.shape == index.shape.

Additionally accepts an optional reduce argument that allows specification of an optional reduction operation, which is applied to all values in the tensor src into input at the indicies specified in the index. For each value in src, the reduction operation is applied to an index in input which is specified by its index in src for dimension != dim and by the corresponding value in index for dimension = dim.

Given a 3-D tensor and reduction using the multiplication operation, input is updated as:

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

Reducing with the addition operation is the same as using oneflow.scatter_add().

The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.Tensor.scatter_.html.

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 or float) – The source blob whose elements will be scatterd and updated to output.

  • reduce (str, optional) – Reduction operation to apply, can be either add or multiply.

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(input, 1, index, src)
>>> out
tensor([[ 0., 10., 20.,  2.,  2.],
        [50., 60.,  2.,  2., 70.],
        [ 2.,  2.,  2.,  2.,  2.]], dtype=oneflow.float32)