torch.sparse.sum

torch.sparse.sum(input, dim=None, dtype=None)

返回給定稀疏張量每一行的總和。

返回稀疏張量 input 在給定維度 dim 上每一行的總和。 如果 dim 是一個維度的列表，則對它們全部進行縮減 (reduce)。 當對所有 sparse_dim 求和時，此方法會返回一個密集張量，而不是稀疏張量。

所有求和的 dim 都會被擠壓 (squeeze) (請參閱 torch.squeeze())，導致輸出張量的維度比 input 少 dim 個維度。

在反向傳播期間，只有 input 的 nnz 位置上的梯度會向後傳播。 請注意，input 的梯度是經過合併的 (coalesced)。

參數

• input (Tensor) – 輸入的稀疏張量

• dim (int 或 int 的 tuple) – 要縮減的維度或維度列表。 預設值：對所有維度進行縮減。

• dtype (torch.dtype, optional) – 返回的 Tensor 的所需數據類型。 預設值：input 的 dtype。

返回類型

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