torch.sparse.mm#

torch.sparse.mm()#

Performs a matrix multiplication of the sparse matrixmat1and the (sparse or strided) matrixmat2. Similar totorch.mm(), ifmat1 is a$(n\times m)$(n×m) tensor,mat2 is a$(m\times p)$(m×p) tensor, out will be a$(n\times p)$(n×p) tensor.Whenmat1 is a COO tensor it must havesparse_dim = 2.When inputs are COO tensors, this function also supports backward for both inputs.

Supports both CSR and COO storage formats.

Note

This function doesn’t support computing derivatives with respect to CSR matrices.

This function also additionally accepts an optionalreduce argument that allowsspecification of an optional reduction operation, mathematically performs the following operation:

$$z_{ij}=\underset{k=0}{\overset{K-1}{⨁}}{x}_{ik}{y}_{kj}$$zij​=k=0⨁K−1​xik​ykj​

where$⨁$⨁ defines the reduce operator.reduce is implemented only forCSR storage format on CPU device.

Parameters

• mat1 (Tensor) – the first sparse matrix to be multiplied

• mat2 (Tensor) – the second matrix to be multiplied, which could be sparse or dense

• reduce (str,optional) – the reduction operation to apply for non-unique indices("sum","mean","amax","amin"). Default"sum".

Shape:

The format of the output tensor of this function follows:- sparse x sparse -> sparse- sparse x dense -> dense

Example:

>>>a=torch.tensor([[1.,0,2],[0,3,0]]).to_sparse().requires_grad_()>>>atensor(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)>>>btensor([[0., 1.],        [2., 0.],        [0., 0.]], requires_grad=True)>>>y=torch.sparse.mm(a,b)>>>ytensor([[0., 1.],        [6., 0.]], grad_fn=<SparseAddmmBackward0>)>>>y.sum().backward()>>>a.gradtensor(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()>>>ctensor(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')>>>y1tensor([[0., 1.],        [6., 0.]], grad_fn=<SparseMmReduceImplBackward0>)>>>y2=torch.sparse.mm(c,b,'max')>>>y2tensor([[0., 1.],        [6., 0.]], grad_fn=<SparseMmReduceImplBackward0>)