gradients

cuda/matmul_cuda.cu

@@ -28,11 +28,11 @@ static void init_cusparse() {
 }
 
 std::tuple<at::Tensor, at::Tensor> spspmm_cuda(at::Tensor A, at::Tensor B) {
+  init_cusparse();
+
   A = A.coalesce();
   B = B.coalesce();
 
-  init_cusparse();
-
   auto m = A.size(0);
   auto n = B.size(1);
   auto k = A.size(1);
@@ -46,6 +46,8 @@ std::tuple<at::Tensor, at::Tensor> spspmm_cuda(at::Tensor A, at::Tensor B) {
   cusparseXcoo2csr(cusparse_handle, indexA[0].data<int>(), nnzA, k,
                    row_ptrA.data<int>(), CUSPARSE_INDEX_BASE_ZERO);
   auto colA = indexA[1];
+  cudaMemcpy(row_ptrA.data<int>() + m, &nnzA, sizeof(int),
+             cudaMemcpyHostToDevice);
 
   auto valueB = B._values();
   auto indexB = B._indices().toType(at::kInt);
@@ -53,6 +55,8 @@ std::tuple<at::Tensor, at::Tensor> spspmm_cuda(at::Tensor A, at::Tensor B) {
   cusparseXcoo2csr(cusparse_handle, indexB[0].data<int>(), nnzB, k,
                    row_ptrB.data<int>(), CUSPARSE_INDEX_BASE_ZERO);
   auto colB = indexB[1];
+  cudaMemcpy(row_ptrB.data<int>() + k, &nnzB, sizeof(int),
+             cudaMemcpyHostToDevice);
 
   cusparseMatDescr_t descr = 0;
   cusparseCreateMatDescr(&descr);

setup.cfg

@@ -5,4 +5,4 @@ description-file = README.md
 test = pytest
 
 [tool:pytest]
-addopts = --capture=no --cov
+addopts = --capture=no

test/test_matmul.py

@@ -3,51 +3,73 @@ from itertools import product
 import pytest
 import torch
 from torch.autograd import gradcheck
-from torch_sparse import spspmm
-from torch_sparse.matmul import SpSpMM
+from torch_sparse import SparseTensor, spspmm, to_value
 
 from .utils import dtypes, devices, tensor
 
+devices = [torch.device('cpu')]
 dtypes = [torch.double]
 
 
 @pytest.mark.parametrize('dtype,device', product(dtypes, devices))
 def test_coalesced_spspmm(dtype, device):
     indexA = torch.tensor([[0, 0, 1, 2, 2], [1, 2, 0, 0, 1]], device=device)
-    valueA = tensor([1, 2, 3, 4, 5], dtype, device, requires_grad=True)
+    valueA = tensor([1, 2, 3, 4, 5], dtype, device)
     sizeA = torch.Size([3, 3])
-    A = (indexA, valueA, sizeA)
-    A_dense = torch.sparse_coo_tensor(indexA, valueA, sizeA).to_dense()
-    A_dense = A_dense.requires_grad_()
-    print('A', A_dense)
+    A = torch.sparse_coo_tensor(indexA, valueA, sizeA, device=device)
 
     indexB = torch.tensor([[0, 2], [1, 0]], device=device)
-    valueB = tensor([2, 4], dtype, device, requires_grad=True)
+    valueB = tensor([2, 4], dtype, device)
     sizeB = torch.Size([3, 2])
-    B = (indexB, valueB, sizeB)
-    B_dense = torch.sparse_coo_tensor(indexB, valueB, sizeB).to_dense()
-    B_dense = B_dense.requires_grad_()
+    B = torch.sparse_coo_tensor(indexB, valueB, sizeB, device=device)
 
-    index, value, size = spspmm(*A, *B)
-    # out = torch.sparse_coo_tensor(index, value, size)
-    expected = torch.matmul(A_dense, B_dense)
-    # assert out.to_dense().tolist() == expected.tolist()
+    assert spspmm(A, B).to_dense().tolist() == [[8, 0], [0, 6], [0, 8]]
 
-    # valueA = valueA.requires_grad_()
-    # valueB = valueB.requires_grad_()
-    # data = (indexA, valueA, sizeA, indexB, valueB, sizeB)
-    # assert gradcheck(SpSpMM.apply, data, eps=1e-6, atol=1e-4) is True
+    # A.requires_grad_()
+    # B.requires_grad_()
 
-    # print(expected)
+    # A.requires_grad_()
+    # B.requires_grad_()
 
-    value.sum().backward()
-    expected.sum().backward()
+    # to_value(C).sum().backward()
+    # print(valueA)
+    # print(valueA.grad)
+    # print(valueB)
+    # print(valueB.grad)
 
-    print(valueA.grad)
-    print(A_dense.grad)
+    # A_dense.requires_grad_()
+    # B_dense.requires_grad_()
 
-    # print(valueB.grad)
+    # C_dense = torch.matmul(A_dense, B_dense)
+    # C_dense[C_dense > 0].sum().backward()
+    # print(A_dense)
+    # print(A_dense.grad)
+    # print(B_dense)
     # print(B_dense.grad)
 
-    # # TODO TEST backward
-    # # value.sum().backward()
+    # A.requires_grad_()
+    # B = B.to_dense()
+    # B.requires_grad_()
+    # torch.spmm(A, B).sum().backward()
+    # print(B.grad)
+
+    # valueA.requires_grad_()
+    valueB.requires_grad_()
+
+    def pipeline(valueA, valueB):
+        A = SparseTensor(indexA, valueA, sizeA)
+        B = SparseTensor(indexB, valueB, sizeB)
+        C = spspmm(A, B)
+        value = to_value(C)
+        return value
+
+    # out = pipeline(valueA, valueB).sum().backward()
+    # print(valueA.grad)
+    # print(valueB.grad)
+
+    print(gradcheck(pipeline, (valueA, valueB), eps=1e-6, atol=1e-4))
+
+    # A, B = Sparsetensor(SparseTensor(index, valueB, sizeB)
+    # print(A.requires_grad)
+
+    # to_value(C).sum().backward()

test/utils.py

@@ -8,5 +8,5 @@ if torch.cuda.is_available():  # pragma: no cover
 
 
 def tensor(x, dtype, device, requires_grad=False):
-    return None if x is None else torch.tensor(
+    return torch.tensor(
         x, dtype=dtype, device=device, requires_grad=requires_grad)

torch_sparse/__init__.py

-from .transpose import transpose
+from .sparse import SparseTensor, to_value
 from .matmul import spspmm
 
 __all__ = [
-    'transpose',
+    'SparseTensor',
+    'to_value',
     'spspmm',
 ]

torch_sparse/matmul.py

 import torch
 import scipy.sparse
-from torch_sparse import transpose
 
 if torch.cuda.is_available():
     import matmul_cuda
 
 
-def spspmm(indexA, valueA, sizeA, indexB, valueB, sizeB):
-    assert valueA.dtype == valueB.dtype
-    assert len(sizeA) == len(sizeB) == 2
-    assert sizeA[1] == sizeB[0]
-
-    index, value = SpSpMM.apply(indexA, valueA, sizeA, indexB, valueB, sizeB)
-    size = torch.Size([sizeA[0], sizeB[1]])
-
-    return index, value, size
-
-
 class SpSpMM(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, indexA, valueA, sizeA, indexB, valueB, sizeB):
-        index, value = mm(indexA, valueA, sizeA, indexB, valueB, sizeB)
-
-        ctx.sizeA, ctx.sizeB = sizeA, sizeB
-        ctx.save_for_backward(indexA, valueA, indexB, valueB, index)
-
-        return index, value
+    def forward(ctx, A, B):
+        ctx.save_for_backward(A, B)
+        return mm(A, B)
 
     @staticmethod
-    def backward(ctx, grad_index, grad_value):
-        indexA, valueA, indexB, valueB, index = ctx.saved_variables
-        grad_valueA = grad_valueB = None
-        grad = (index, grad_value, torch.Size([ctx.sizeA[0], ctx.sizeB[1]]))
+    def backward(ctx, grad_C):
+        A, B = ctx.saved_variables
+        grad_A = grad_B = None
+
+        if ctx.needs_input_grad[0]:
+            grad_A = mm(grad_C, B.t().coalesce())
 
         if ctx.needs_input_grad[1]:
-            B_tranposed = transpose(indexB, valueB, ctx.sizeB)
-            _, grad_valueA = mm(*grad, *B_tranposed)
+            grad_B = mm(A.t(), grad_C)
 
-        if ctx.needs_input_grad[4]:
-            A_tranposed = transpose(indexA, valueA, ctx.sizeA)
-            _, grad_valueB = mm(*A_tranposed, *grad)
+        return grad_A, grad_B
 
-        return None, grad_valueA, None, None, grad_valueB, None
 
+spspmm = SpSpMM.apply
 
-def mm(indexA, valueA, sizeA, indexB, valueB, sizeB):
-    if valueA.is_cuda:
-        return mm_cuda(indexA, valueA, sizeA, indexB, valueB, sizeB)
-    else:
-        return mm_cpu(indexA, valueA, sizeA, indexB, valueB, sizeB)
 
+def mm(A, B):
+    assert A.dtype == B.dtype
+    assert A.size(1) == B.size(0)
+    return mm_cuda(A, B) if A.is_cuda else mm_cpu(A, B)
 
-def mm_cuda(indexA, valueA, sizeA, indexB, valueB, sizeB):
-    A = torch.sparse_coo_tensor(indexA, valueA, sizeA)
-    B = torch.sparse_coo_tensor(indexB, valueB, sizeB)
 
+def mm_cuda(A, B):
     index, value = matmul_cuda.spspmm(A, B)
-
-    return index, value
+    size = torch.Size([A.size(0), B.size(1)])
+    return torch.sparse_coo_tensor(index, value, size, device=value.device)
 
 
-def mm_cpu(indexA, valueA, sizeA, indexB, valueB, sizeB):
-    A, B, = to_scipy(indexA, valueA, sizeA), to_scipy(indexB, valueB, sizeB)
-    C = A.tocsr().dot(B.tocsr()).tocoo()
+def mm_cpu(A, B):
+    return from_scipy(to_scipy(A).dot(to_scipy(B)))
 
-    row, col = torch.from_numpy(C.row).long(), torch.from_numpy(C.col).long()
-    index = torch.stack([row, col], dim=0)
-    value = torch.from_numpy(C.data).type_as(valueA)
 
-    return index, value
+def to_scipy(A):
+    (row, col), data, shape = A._indices(), A._values(), tuple(A.size())
+    row, col, data = row.detach(), col.detach(), data.detach()
+    return scipy.sparse.coo_matrix((data, (row, col)), shape).tocsr()
 
 
-def to_scipy(index, value, size):
-    (row, col), value = index.detach().numpy(), value.detach().numpy()
-    return scipy.sparse.coo_matrix((value, (row, col)), tuple(size))
+def from_scipy(A):
+    A = A.tocoo()
+    row, col, value, size = A.row, A.col, A.data, torch.Size(A.shape)
+    value = torch.from_numpy(value)
+    index = torch.stack([torch.from_numpy(row), torch.from_numpy(col)], dim=0)
+    index = index.to(torch.long)
+    return torch.sparse_coo_tensor(index, value, size)

torch_sparse/sparse.py

+import torch
+
+
+class _SparseTensor(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, index, value, size):
+        ctx.size = size
+        ctx.save_for_backward(index)
+        return torch.sparse_coo_tensor(index, value, size, device=value.device)
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        index = ctx.saved_variables[0]
+        grad_in = None
+
+        if ctx.needs_input_grad[1]:
+            value = grad_out._values()
+            id1 = index[0] * ctx.size[1] + index[1]
+            index = grad_out._indices()
+            id2 = index[0] * ctx.size[1] + index[1]
+
+            grad_in = value.new_zeros(id1.max().item() + 1)
+            grad_in[id2] = value
+            grad_in = grad_in[id1]
+
+        return None, grad_in, None
+
+
+SparseTensor = _SparseTensor.apply
+
+
+class ToValue(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, A):
+        ctx.save_for_backward(A)
+        return A._values()
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        A = ctx.saved_variables[0]
+        grad_in = None
+
+        if ctx.needs_input_grad[0]:
+            grad_in = torch.sparse_coo_tensor(
+                A._indices(), grad_out, A.size(), device=grad_out.device)
+
+        return grad_in
+
+
+to_value = ToValue.apply

torch_sparse/transpose.py

-import torch
-
-
-def transpose(index, value, size):
-    (row, col), (dim1, dim2) = index, size
-    index, size = torch.stack([col, row], dim=0), torch.Size([dim2, dim1])
-    return index, value, size