torch.Tensor.index_add_

Tensor.index_add_(dim, index, source, *, alpha=1) → Tensor

Накапливает элементы alpha, умноженные на source, в тензоре self, добавляя их к индексам в порядке, заданном в index. Например, если dim == 0, index[i] == j и alpha=-1, то i-я строка source вычитается из j-й строки self.

Размер 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. Дополнительные сведения см. в разделе Воспроизводимость.

Параметры:

• dim (int) – измерение, вдоль которого выполняется индексирование
• index (Tensor) – индексы source для выбора; тип данных должен быть либо torch.int64, либо torch.int32
• source (Tensor) – тензор, содержащий добавляемые значения

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

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