torch.Tensor.scatter_reduce_

Tensor.scatter_reduce_(dim, index, src, reduce, *, include_self=True) → Tensor

使用通过 reduce 参数定义的应用归约（"sum"、"prod"、"mean"、"amax"、"amin"），将 src 张量中的所有值归约到 self 张量中由 index 张量指定的索引。对于 src 中的每个值，它都会被归约到 self 中的一个索引，该索引由 src 中 dimension != dim 的索引和 dimension = dim 的 index 中的相应值指定。如果 include_self="True"，则 self 张量中的值将包含在归约中。

self、index 和 src 应具有相同的维度数。还需要 index.size(d) <= src.size(d) 对于所有维度 d，以及 index.size(d) <= self.size(d) 对于所有维度 d != dim。请注意，index 和 src 不广播。

对于 reduce="sum" 和 include_self=True 的 3-D 张量，输出给定为

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

注意

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

注意

仅针对 src.shape == index.shape 实现了反向传递。

警告

此功能处于 Beta 版，在不久的将来可能会发生更改。

参数

• dim (int) – 沿其索引的轴

• index (LongTensor) – 要分散和归约的元素的索引。

• src (Tensor) – 要分散和归约的源元素

• reduce (str) – 应用于非唯一索引的归约操作 ("sum", "prod", "mean", "amax", "amin")

• include_self (bool) – self 张量中的元素是否包含在归约中

示例

>>> src = torch.tensor([1., 2., 3., 4., 5., 6.])
>>> index = torch.tensor([0, 1, 0, 1, 2, 1])
>>> input = torch.tensor([1., 2., 3., 4.])
>>> input.scatter_reduce(0, index, src, reduce="sum")
tensor([5., 14., 8., 4.])
>>> input.scatter_reduce(0, index, src, reduce="sum", include_self=False)
tensor([4., 12., 5., 4.])
>>> input2 = torch.tensor([5., 4., 3., 2.])
>>> input2.scatter_reduce(0, index, src, reduce="amax")
tensor([5., 6., 5., 2.])
>>> input2.scatter_reduce(0, index, src, reduce="amax", include_self=False)
tensor([3., 6., 5., 2.])