torch.Tensor.index_reduce_

Tensor.index_reduce_(dim, index, source, reduce, *, include_self=True) → Tensor

Накопить элементы source в тензоре self, накапливая их по индексам в порядке, заданном в index, с помощью операции редукции, указанной в аргументе reduce. Например, если dim == 0, index[i] == j, reduce == prod и include_self == True, то строка i-го измерения source умножается на строку j-го измерения self. Если include_self="True", значения тензора self включаются в редукцию; в противном случае строки тензора self, в которые выполняется накопление, рассматриваются так, как если бы они были заполнены нейтральными элементами редукции.

Измерение dim source должно иметь размер, равный длине index (который должен быть вектором), а размеры всех остальных измерений должны совпадать с self, иначе будет вызвана ошибка.

Для 3D-тензора с reduce="prod" и include_self=True результат задается следующим образом:

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 эта операция может быть недетерминированной. Дополнительную информацию см. в разделе Воспроизводимость.

Примечание

Эта функция поддерживает только тензоры с числами с плавающей точкой.

Предупреждение

Эта функция находится в бета-версии и может измениться в ближайшем будущем.

Параметры:

• dim (int) – измерение, по которому выполняется индексация
• index (Tensor) – индексы self, в которые выполняется накопление; dtype должен быть torch.int64 или torch.int32
• source (FloatTensor) – тензор со значениями для накопления
• reduce (str) – применяемая операция редукции ("prod", "mean", "amax", "amin")

Именованные аргументы:

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.]])