On this page
torch.Tensor.index_add_
Tensor.index_add_(dim, index, source, *, alpha=1) → Tensor
-
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
Note
This operation may behave nondeterministically when given tensors on a CUDA device. See Reproducibility for more information.
- Parameters
- Keyword Arguments
-
alpha (Number) – the scalar multiplier for
source
Example:
>>> x = torch.ones(5, 3) >>> t = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float) >>> index = torch.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.]]) >>> x.index_add_(0, index, t, alpha=-1) tensor([[ 1., 1., 1.], [ 1., 1., 1.], [ 1., 1., 1.], [ 1., 1., 1.], [ 1., 1., 1.]])
© 2024, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://pytorch.org/docs/2.1/generated/torch.Tensor.index_add_.html