torch.Tensor.index_add_¶
- Tensor.index_add_(dim, index, source, *, alpha=1) Tensor ¶
通过按照
index
中给定的顺序添加到索引,将alpha
倍的source
的元素累加到self
张量中。例如,如果dim == 0
,index[i] == j
且alpha=-1
,则source
的第i
行将从self
的第j
行中减去。source
的第dim
维的大小必须与index
(它必须是向量)的长度相同,并且所有其他维度必须与self
匹配,否则将引发错误。对于 3 维张量,输出如下所示
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 设备上给出张量时,此操作可能会表现出不确定性行为。有关更多信息,请参阅 可重复性。
- 参数
- 关键字参数
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.]])