def test_cpu_matmul(self, mat1, mat2, expected):
out = torch.empty_like(expected)
mat1_csr = SparseTensor.from_scipy(scipy.sparse.csr_matrix(mat1))
mat2_csc = SparseTensor.from_scipy(scipy.sparse.csc_matrix(mat2))
sparse_matmul(mat1_csr, mat2_csc, out)
torch.testing.assert_allclose(out, expected)
def test_cuda_matmul(self, mat1, mat2, expected):
dev = torch.device("cuda:0")
out = create_fortran(expected.shape, expected.dtype, dev)
mat1_csr = SparseTensor.from_scipy(
scipy.sparse.csr_matrix(mat1)).to(device=dev)
mat2_csr = SparseTensor.from_scipy(
scipy.sparse.csr_matrix(mat2)).to(device=dev)
sparse_matmul(mat1_csr, mat2_csr, out)
torch.testing.assert_allclose(out.cpu(), expected)
def test_cpu_matmul_wrong_format(self, mat1, mat2, expected):
out = torch.empty_like(expected)
mat1_csr = SparseTensor.from_scipy(scipy.sparse.csr_matrix(mat1))
mat2_csr = SparseTensor.from_scipy(scipy.sparse.csr_matrix(mat2))
with pytest.raises(ValueError) as exc_info:
sparse_matmul(mat1_csr, mat2_csr, out)
assert str(exc_info.value).startswith("B must be CSC matrix")
mat1_csc = SparseTensor.from_scipy(scipy.sparse.csc_matrix(mat1))
with pytest.raises(ValueError) as exc_info:
sparse_matmul(mat1_csc, mat2_csr, out)
assert str(exc_info.value).startswith("A must be CSR matrix")
def test_matmul_zeros(self, mat1, mat2, expected, device):
mat1_zero_csr = SparseTensor.from_scipy(
scipy.sparse.csr_matrix(
torch.zeros_like(mat1).numpy())).to(device=device)
mat2_csc = SparseTensor.from_scipy(
scipy.sparse.csc_matrix(mat2.numpy())).to(device=device)
out = torch.empty_like(expected).to(device)
sparse_matmul(mat1_zero_csr, mat2_csc, out)
assert torch.all(out == 0.0)
mat1_csr = SparseTensor.from_scipy(
scipy.sparse.csr_matrix(mat1.numpy())).to(device=device)
mat2_zero_csc = SparseTensor.from_scipy(
scipy.sparse.csc_matrix(
torch.zeros_like(mat2).numpy())).to(device=device)
out = torch.empty_like(expected).to(device=device)
sparse_matmul(mat1_csr, mat2_zero_csc, out)
assert torch.all(out == 0.0)