torch.sparse.mm

torch.sparse.mm()

执行稀疏矩阵mat1和（稀疏或步进）矩阵mat2之间的矩阵乘法。类似于torch.mm()，如果mat1是一个 $(n \times m)$ 张量，而mat2是一个 $(m \times p)$ 张量，则输出out将是一个 $(n \times p)$ 张量。当mat1是COO张量时，它必须具有sparse_dim = 2属性。如果输入为COO张量，则此函数还支持对两个输入的反向传播。

支持_CSR_和_COO_存储格式。

注意

此函数不支持计算与CSR矩阵相关的导数。

此函数还接受一个可选的 reduce 参数，用于指定可选的缩减操作，其数学运算如下：

z_{ij} = \bigoplus_{k = 0}^{K - 1} x_{ik} y_{kj}

其中 $\bigoplus$ 表示减少操作符。函数reduce仅在CPU设备上的CSR存储格式中被实现。

参数

mat1 (Tensor) – 第一个要相乘的稀疏矩阵
mat2 (Tensor) – 要与第一个矩阵相乘的第二个矩阵，可以是稀疏矩阵或稠密矩阵
reduce (str, 可选) – 对非唯一索引应用的缩减操作（"sum"，"mean"，"amax"，"amin"）。默认值为 "sum"。

形状:: 该函数输出张量的格式如下：稀疏矩阵乘以稀疏矩阵得到稀疏结果，稀疏矩阵乘以密集矩阵得到密集结果。

示例：

>>> a = torch.tensor([[1., 0, 2], [0, 3, 0]]).to_sparse().requires_grad_()
>>> a
tensor(indices=tensor([[0, 0, 1],
                       [0, 2, 1]]),
       values=tensor([1., 2., 3.]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo, requires_grad=True)
>>> b = torch.tensor([[0, 1.], [2, 0], [0, 0]], requires_grad=True)
>>> b
tensor([[0., 1.],
        [2., 0.],
        [0., 0.]], requires_grad=True)
>>> y = torch.sparse.mm(a, b)
>>> y
tensor([[0., 1.],
        [6., 0.]], grad_fn=<SparseAddmmBackward0>)
>>> y.sum().backward()
>>> a.grad
tensor(indices=tensor([[0, 0, 1],
                       [0, 2, 1]]),
       values=tensor([1., 0., 2.]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)
>>> c = a.detach().to_sparse_csr()
>>> c
tensor(crow_indices=tensor([0, 2, 3]),
       col_indices=tensor([0, 2, 1]),
       values=tensor([1., 2., 3.]), size=(2, 3), nnz=3,
       layout=torch.sparse_csr)
>>> y1 = torch.sparse.mm(c, b, 'sum')
>>> y1
tensor([[0., 1.],
        [6., 0.]], grad_fn=<SparseMmReduceImplBackward0>)
>>> y2 = torch.sparse.mm(c, b, 'max')
>>> y2
tensor([[0., 1.],
        [6., 0.]], grad_fn=<SparseMmReduceImplBackward0>)