torch.Tensor.scatter_add_¶
- Tensor.scatter_add_(dim, index, src) Tensor ¶
将张量
src
中的所有值添加到self
中,索引在index
张量中指定,类似于scatter_()
。对于src
中的每个值,它将被添加到self
中的索引,该索引由src
中的索引指定,用于dimension != dim
,并由index
中的对应值指定,用于dimension = dim
。对于一个 3 维张量,
self
将更新为self[index[i][j][k]][j][k] += src[i][j][k] # if dim == 0 self[i][index[i][j][k]][k] += src[i][j][k] # if dim == 1 self[i][j][index[i][j][k]] += src[i][j][k] # if dim == 2
self
、index
和src
应具有相同的维度数量。还需要index.size(d) <= src.size(d)
对于所有维度d
,并且index.size(d) <= self.size(d)
对于所有维度d != dim
。请注意,index
和src
不进行广播。注意
此操作在 CUDA 设备上给出张量时可能表现出非确定性行为。有关详细信息,请参阅 可重复性。
注意
反向传播仅在
src.shape == index.shape
时实现。- 参数
示例
>>> src = torch.ones((2, 5)) >>> index = torch.tensor([[0, 1, 2, 0, 0]]) >>> torch.zeros(3, 5, dtype=src.dtype).scatter_add_(0, index, src) tensor([[1., 0., 0., 1., 1.], [0., 1., 0., 0., 0.], [0., 0., 1., 0., 0.]]) >>> index = torch.tensor([[0, 1, 2, 0, 0], [0, 1, 2, 2, 2]]) >>> torch.zeros(3, 5, dtype=src.dtype).scatter_add_(0, index, src) tensor([[2., 0., 0., 1., 1.], [0., 2., 0., 0., 0.], [0., 0., 2., 1., 1.]])