torch.sparse.sum

torch.sparse.sum(input, dim=None, dtype=None) [исходный код]

Возвращает сумму элементов каждой строки заданного разреженного тензора.

Возвращает сумму элементов каждой строки разреженного тензора input по заданным измерениям dim. Если dim — это список измерений, выполняется редукция по всем им. При суммировании по всем sparse_dim этот метод возвращает плотный тензор вместо разреженного.

Все суммированные dim удаляются (см. torch.squeeze()), поэтому выходной тензор имеет на dim меньше измерений, чем input.

При обратном распространении градиенты будут распространяться только в nnz позициях input. Обратите внимание, что градиенты input объединены.

Параметры:

• input (Tensor) – входной разреженный тензор
• dim (int или tuple из ints) – измерение или список измерений для редукции. По умолчанию: редукция по всем измерениям.
• dtype (torch.dtype, необязательно) – требуемый тип данных возвращаемого тензора. По умолчанию: тип данных input.

Тип возвращаемого значения:

Tensor

Пример:

>>> nnz = 3
>>> dims = [5, 5, 2, 3]
>>> I = torch.cat([torch.randint(0, dims[0], size=(nnz,)),
                   torch.randint(0, dims[1], size=(nnz,))], 0).reshape(2, nnz)
>>> V = torch.randn(nnz, dims[2], dims[3])
>>> size = torch.Size(dims)
>>> with torch.sparse.check_sparse_tensor_invariants():
...     S = torch.sparse_coo_tensor(I, V, size)
>>> S
tensor(indices=tensor([[2, 0, 3],
                       [2, 4, 1]]),
       values=tensor([[[-0.6438, -1.6467,  1.4004],
                       [ 0.3411,  0.0918, -0.2312]],

                      [[ 0.5348,  0.0634, -2.0494],
                       [-0.7125, -1.0646,  2.1844]],

                      [[ 0.1276,  0.1874, -0.6334],
                       [-1.9682, -0.5340,  0.7483]]]),
       size=(5, 5, 2, 3), nnz=3, layout=torch.sparse_coo)

# when sum over only part of sparse_dims, return a sparse tensor
>>> torch.sparse.sum(S, [1, 3])
tensor(indices=tensor([[0, 2, 3]]),
       values=tensor([[-1.4512,  0.4073],
                      [-0.8901,  0.2017],
                      [-0.3183, -1.7539]]),
       size=(5, 2), nnz=3, layout=torch.sparse_coo)

# when sum over all sparse dim, return a dense tensor
# with summed dims squeezed
>>> torch.sparse.sum(S, [0, 1, 3])
tensor([-2.6596, -1.1450])