Shortcuts

torch.Tensor.index_add_

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

alpha 乘以 source 的元素累加到 self 张量中,按照 index 中给定的顺序添加到索引中。例如,如果 dim == 0index[i] == j,并且 alpha=-1,那么 source 的第 i 行将从 self 的第 j 行中减去。

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

对于一个3维张量,输出如下所示:

self[index[i], :, :] += alpha * src[i, :, :]  # 如果 dim == 0
self[:, index[i], :] += alpha * src[:, i, :]  # 如果 dim == 1
self[:, :, index[i]] += alpha * src[:, :, i]  # 如果 dim == 2

注意

当在CUDA设备上使用张量时,此操作可能会表现出不确定性行为。更多信息请参见可重复性

Parameters
  • dim (int) – 沿此维度进行索引

  • 索引 (张量) – 用于从 source 中选择的索引, 应具有 torch.int64torch.int32 的数据类型

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

Keyword Arguments

alpha (数字) – 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.]])
优云智算