torch.Tensor.index_add_ — PyTorch 2.7 documentation (original) (raw)

Tensor.index_add_(dim, index, source, *, alpha=1) → Tensor¶

Accumulate the elements of alpha times source into the selftensor by adding to the indices in the order given in index. For example, if dim == 0, index[i] == j, and alpha=-1, then the ith row ofsource is subtracted from the jth row of self.

The dimth dimension of source must have the same size as the length of index (which must be a vector), and all other dimensions must match self, 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

• dim (int) – dimension along which to index
• index (Tensor) – indices of source to select from, should have dtype either torch.int64 or torch.int32
• source (Tensor) – the tensor containing values to add

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.]])