torch.Tensor.index_reduce_¶
- Tensor.index_reduce_(dim, index, source, reduce, *, include_self=True) Tensor ¶
使用
reduce
参数给出的归约,通过按index
中给出的顺序累加到索引来将source
的元素累加到self
张量中。例如,如果dim == 0
、index[i] == j
、reduce == prod
且include_self == True
,则source
的第i
行将乘以self
的第j
行。如果include_self="True"
,则self
张量中的值将包含在归约中,否则,累加到的self
张量中的行将被视为用归约标识符填充。source
的第dim
维必须与index
的长度(必须是向量)相同,并且所有其他维必须与self
相匹配,否则将引发错误。对于具有
reduce="prod"
和include_self=True
的 3-D 张量,输出将给出为self[index[i], :, :] *= src[i, :, :] # if dim == 0 self[:, index[i], :] *= src[:, i, :] # if dim == 1 self[:, :, index[i]] *= src[:, :, i] # if dim == 2
注意
此操作在 CUDA 设备上给出张量时可能表现出不确定性。有关更多信息,请参阅 可重复性。
注意
此函数仅支持浮点张量。
警告
此函数处于测试阶段,未来可能会更改。
- 参数
- 关键字参数
include_self (bool) –
self
张量中的元素是否包含在归约中
示例
>>> x = torch.empty(5, 3).fill_(2) >>> t = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=torch.float) >>> index = torch.tensor([0, 4, 2, 0]) >>> x.index_reduce_(0, index, t, 'prod') tensor([[20., 44., 72.], [ 2., 2., 2.], [14., 16., 18.], [ 2., 2., 2.], [ 8., 10., 12.]]) >>> x = torch.empty(5, 3).fill_(2) >>> x.index_reduce_(0, index, t, 'prod', include_self=False) tensor([[10., 22., 36.], [ 2., 2., 2.], [ 7., 8., 9.], [ 2., 2., 2.], [ 4., 5., 6.]])