spspmm gives wrong output

[gao462]

The sparse-sparse multiplication (spspmm) gives wrong output which is different from SciPy sparse matrix multiplication.

import numpy as onp
import torch
import scipy.sparse as sparse
from torch_sparse import spspmm

num_srcs = 3
num_dsts = 4
rows = [0, 1, 2]
cols = [3, 2, 1]
vals = [1.0, 1.0, 1.0]
spmat = sparse.coo_matrix((vals, (rows, cols)), shape=(num_srcs, num_dsts))
print("SciPy Sparse")
print((spmat @ spmat.T).todense())

print("Dense")
dsmat = spmat.todense()
print(dsmat @ dsmat.T)

indices = torch.tensor([rows, cols])
values = torch.tensor(vals)
(indices, values) = spspmm(indices, values, indices[[1, 0]], values, num_srcs, num_dsts, num_srcs)
res = torch.zeros(num_srcs, num_srcs)
res[indices[0], indices[1]] = values
print("Torch Sparse")
print(res.numpy())

Get results:

SciPy Sparse
[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]
Dense
[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]
Torch Sparse
[[0. 0. 0.]
 [0. 0. 0.]
 [1. 1. 1.]]

torch_sparse version is 0.6.14.

[rusty1s]

It is best to use the SparseTensor class directly for this:

from torch_sparse import SparseTensor

spmat = SparseTensor(row=torch.from_numpy(rows),
                     col=torch.from_numpy(cols),
                     value=torch.from_numpy(vals),
                     sparse_sizes=(3, 4))
print(spmat @ spmat.t())

The spspmm interface is pretty old and not well maintained. You can fix your initial error by running

(indices, values) = spspmm(indices, values, indices[[1, 0]], values, num_srcs,
                           num_dsts, num_srcs, coalesced=True)