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 tensorsrc
intoinput
at the indicies specified in theindex
. For each value insrc
, the reduction operation is applied to an index ininput
which is specified by its index insrc
fordimension != dim
and by the corresponding value inindex
fordimension = 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
ormultiply
.
- Returns
The scatterd Tensor.
- Return type
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)