oneflow.Tensor.index_add_¶
-
Tensor.
index_add_
(dim, index, source, *, alpha=1) → Tensor¶ The interface is consistent with PyTorch.
Accumulate the elements of
alpha
timessource
into theself
tensor by adding to the indices in the order given inindex
. For example, ifdim == 0
,index[i] == j
, andalpha=-1
, then thei
th row ofsource
is subtracted from thej
th row ofself
.The
dim
th dimension ofsource
must have the same size as the length ofindex
(which must be a vector), and all other dimensions must matchself
, or an error will be raised.For a 3-D tensor the output is given as:
self[index[i], :, :] += alpha * src[i, :, :] # if dim == 0 self[:, index[i], :] += alpha * src[:, i, :] # if dim == 1 self[:, :, index[i]] += alpha * src[:, :, i] # if dim == 2
- Parameters
- Keyword Arguments
alpha (Number) – the scalar multiplier for
source
>>> import oneflow as flow >>> x = flow.ones(5, 3) >>> t = flow.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=flow.float) >>> index = flow.tensor([0, 4, 2]) >>> x.index_add_(0, index, t) tensor([[ 2., 3., 4.], [ 1., 1., 1.], [ 8., 9., 10.], [ 1., 1., 1.], [ 5., 6., 7.]], dtype=oneflow.float32) >>> x.index_add_(0, index, t, alpha=-1) tensor([[1., 1., 1.], [1., 1., 1.], [1., 1., 1.], [1., 1., 1.], [1., 1., 1.]], dtype=oneflow.float32)