torch.Tensor.index_add_

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

将 alpha 的元素累加到 self 张量中，累加次数为 source 乘以 index 中给出的索引顺序。例如，如果 dim == 0 ， index[i] == j ，和 alpha=-1 ，那么 i 行 source 的值从 j 行 self 的值中减去。

第 dim 维度 source 必须与 index 的长度相同（index 必须是一个向量），并且所有其他维度必须与 self 匹配，否则将引发错误。

对于三维张量，输出如下：

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

注意

当在 CUDA 设备上给定张量时，此操作可能表现出非确定性。有关更多信息，请参阅可重现性。

参数:

• dim（int）- 沿着哪个维度进行索引

• index (Tensor) – 从 source 选择的索引，应具有 torch.int64 或 torch.int32 的数据类型

• source (Tensor) – 包含要添加值的张量

关键字参数:

alpha (Number) – source 的标量乘数

示例：

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