torch.Tensor.index_add_

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

將 alpha 乘以 source 的元素累加到 self 張量中，透過按 index 中給定順序新增到索引位置。例如，如果 dim == 0， index[i] == j，且 alpha=-1，則 source 的第 i 行將從 self 的第 j 行中減去。

source 的第 dim 維度必須與 index 的長度（index 必須是向量）具有相同的大小，所有其他維度必須與 self 匹配，否則將引發錯誤。

對於 3-D 張量，輸出如下所示：

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 中選擇的索引，dtype 應為 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.]])