torch.sparse.mm¶
- torch.sparse.mm()¶
执行稀疏矩阵
mat1
和(稀疏或带步长)矩阵mat2
的矩阵乘法。类似于torch.mm()
,如果mat1
是一个 张量,mat2
是一个 张量,输出将是一个 张量。当mat1
是一个 COO 张量时,它必须具有 sparse_dim = 2。当输入是 COO 张量时,此函数还支持对两个输入的 backward 运算。支持 CSR 和 COO 存储格式。
注意
此函数不支持计算相对于 CSR 矩阵的导数。
此函数还另外接受一个可选的
reduce
参数,该参数允许指定可选的约简操作,在数学上执行以下操作其中 定义了约简运算符。
reduce
仅在 CPU 设备上针对 CSR 存储格式实现。- 参数
- 形状
此函数的输出张量的格式如下: - 稀疏 x 稀疏 -> 稀疏 - 稀疏 x 稠密 -> 稠密
示例
>>> 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>)