Repository: cai4cai/torchsparsegradutils (language: python)

Issue:
Add memory-efficient sparse @ dense op
CSR matrix multiplication in PyTorch is not currently backpropagatable. There is a fix but it is post 1.13 branch:
https://github.com/pytorch/pytorch/pull/86301

Sparse matrix multiplication with sparse gradient is not memory efficient:
https://github.com/pytorch/pytorch/issues/41128

A workaround for this and the previous point can be done similar to
https://github.com/pytorch/pytorch/issues/41128#issuecomment-1286106119


Discussion / hints:

Fix patch:
diff --git a/torchsparsegradutils/__init__.py b/torchsparsegradutils/__init__.py
index e416a76..909cd5b 100644
--- a/torchsparsegradutils/__init__.py
+++ b/torchsparsegradutils/__init__.py
@@ -1,3 +1,4 @@
+from .sparse_matmul import sparse_mm
+from .sparse_solve import sparse_triangular_solve
 
-def dummy_interface():
-    print("called dummy_interface() function")
\ No newline at end of file
+__all__ = ['sparse_mm', 'sparse_triangular_solve']
\ No newline at end of file
diff --git a/torchsparsegradutils/sparse_matmul.py b/torchsparsegradutils/sparse_matmul.py
new file mode 100644
index 0000000..1e3e5c5
--- /dev/null
+++ b/torchsparsegradutils/sparse_matmul.py
@@ -0,0 +1,65 @@
+import torch
+
+def sparse_mm(A, B):
+    return SparseMatMul.apply(A, B)
+
+class SparseMatMul(torch.autograd.Function):
+    """
+    Matrix multuplication between sparse matrix (A) 
+    and dense matrix (B), with support for backpropagation
+    Matrix A can be in either COO or CSR format
+    
+    This implementation provides a memory efficient version of
+    torch.sparse.mm on the backward pass, working around the issue described in:
+    https://github.com/pytorch/pytorch/issues/41128
+    """
+    @staticmethod
+    def forward(ctx, A, B):
+        grad_flag = A.requires_grad or B.requires_grad
+        A, B = A.detach(), B.detach()
+        x = torch.sparse.mm(A, B)
+        x.requires_grad = grad_flag
+        ctx.save_for_backward(A, B)
+        return x
+
+    @staticmethod
+    def backward(ctx, grad):
+        A, B = ctx.saved_tensors
+        
+        # The gradient with respect to the matrix A, seen as a dense matrix, would
+        # lead to a backprop rule as follows: gradA = grad @ b.T
+        # but we are only interested in the gradient with respect to
+        # the (non-zero) values of A. To save memory, instead of computing the full
+        # dense matrix prev_grad @ b and then subsampling at the nnz locations in A,
+        # we can directly only compute the required values:
+        # grad_a[i,j] = dotprod(grad[i,:], b[j,:])
+        
+        # We start by getting the i and j indices:
+        
+        if A.layout == torch.sparse_coo:
+            A_row_idx, A_col_idx = A.indices()
+        elif A.layout == torch.sparse_csr:
+            A_col_idx = A.col_indices()
+            A_crow_idx = A.crow_indices()
+            # Uncompress row indices:
+            A_row_idx = torch.repeat_interleave(torch.arange(A.size()[0], device=A.device), 
+                                                A_crow_idx[1:]-A_crow_idx[:-1])
+        else:
+            raise ValueError(f"Unsupported layout: {A.layout}")
+
+        grad_select = grad.index_select(0, A_row_idx)  # grad[i, :]
+        B_select = B.index_select(0, A_col_idx)  # B[j, :]
+        
+        # Dot product: 
+        gradB_ewise = grad_select * B_select
+        gradA = torch.sum(gradB_ewise, dim=1)
+        
+        # Create a sparse matrix of the gradient with respect to the nnz of A
+        if A.layout == torch.sparse_coo:
+            gradA = torch.sparse_coo_tensor(A.indices(), gradA, A.shape)
+        elif A.layout == torch.sparse_csr:
+            gradA = torch.sparse_csr_tensor(A_crow_idx, A_col_idx, gradA, A.shape)
+        
+        # Now compute the dense gradient with respect to B
+        gradB = torch.sparse.mm(A.t(), grad)
+        return gradA, gradB
\ No newline at end of file
diff --git a/torchsparsegradutils/sparse_solve.py b/torchsparsegradutils/sparse_solve.py
new file mode 100644
index 0000000..70e91d7
--- /dev/null
+++ b/torchsparsegradutils/sparse_solve.py
@@ -0,0 +1,70 @@
+import torch
+
+def sparse_triangular_solve(A, B, upper=True):
+    return SparseTriangularSolve.apply(A, B, upper)
+    
+
+class SparseTriangularSolve(torch.autograd.Function):
+    """
+    Solves a system of equations with a square upper or lower triangular 
+    invertible sparse matrix A and dense right-hand side matrix B, 
+    with backpropagation support
+    
+    Solves: Ax = B
+
+    A can be in either COO or CSR format.
+    But, COO will internally be converted to CSR before solving.
+    
+    This implementation preserves the sparsity of the gradient calculated during a 
+    backpass through torch.triangular_solve, as detailed here:
+    https://github.com/pytorch/pytorch/issues/87358
+    """
+    @staticmethod
+    def forward(ctx, A, B, upper):
+        grad_flag = A.requires_grad or B.requires_grad
+        ctx.csr = True
+        ctx.upper = upper
+        if A.layout == torch.sparse_coo:  
+            A = A.to_sparse_csr() # triangular solve doesn't work with sparse coo
+            ctx.csr = False      
+        x = torch.triangular_solve(B.detach(), A.detach(), upper=upper).solution
+        x.requires_grad = grad_flag
+        ctx.save_for_backward(A, x.detach())
+        return x
+
+    @staticmethod
+    def backward(ctx, grad):
+        A, x = ctx.saved_tensors
+        
+        # Backprop rule: gradB = a^{-T} grad
+        gradB = torch.triangular_solve(grad, A, upper=ctx.upper, transpose=True).solution
+        
+        # The gradient with respect to the matrix a seen as a dense matrix would
+        # lead to a backprop rule as follows
+        # gradA = -(A^{-T} grad)(A^{-1} B) = - gradB @ x.T
+        # but we are only interested in the gradient with respect to
+        # the (non-zero) values of A. To save memory, instead of computing the full
+        # dense matrix gradB @ x.T and then subsampling at the nnz locations in a,
+        # we can directly only compute the required values:
+        # gradA[i,j] = - dotprod(gradB[i,:], x[j,:])
+        
+        # We start by getting the i and j indices:
+        A_col_idx = A.col_indices()
+        A_crow_idx = A.crow_indices()
+        # Uncompress row indices:
+        A_row_idx = torch.repeat_interleave(torch.arange(A.size()[0], device=A.device), 
+                                            A_crow_idx[1:]-A_crow_idx[:-1])     
+        
+        mgradbselect = -gradB.index_select(0, A_row_idx)   # -gradB[i, :]
+        xselect = x.index_select(0, A_col_idx)  # x[j, :]
+        
+        # Dot product:
+        mgbx = mgradbselect * xselect
+        gradA = torch.sum(mgbx, dim=1)
+            
+        if ctx.csr is False:
+            gradA = torch.sparse_coo_tensor(torch.stack([A_row_idx, A_col_idx]), gradA, A.shape)
+        else:
+            gradA = torch.sparse_csr_tensor(A_crow_idx, A_col_idx, gradA, A.shape)
+
+        return gradA, gradB, None
\ No newline at end of file


Test patch:
diff --git a/tests/test_sparse_matmul.py b/tests/test_sparse_matmul.py
new file mode 100644
index 0000000..854e77d
--- /dev/null
+++ b/tests/test_sparse_matmul.py
@@ -0,0 +1,95 @@
+import torch
+import unittest
+from random import randrange
+from torchsparsegradutils import sparse_mm
+
+
+def gencoordinates(nr, nc, ni, device='cuda'):
+    """Used to genererate ni random unique coordinates for sparse matrix with size [nr, nc]"""
+    coordinates = set()
+    while True:
+        r, c = randrange(nr), randrange(nc)
+        coordinates.add((r, c))
+        if len(coordinates) == ni:
+            return torch.stack([torch.tensor(co) for co in coordinates], dim=-1).to(device)
+
+
+class SparseMatMulTest(unittest.TestCase):
+    """Test Sparse x Dense matrix multiplication with back propagation for COO and CSR matrices"""
+    def setUp(self) -> None:
+        self.A_shape = (8, 16)
+        self.B_shape = (self.A_shape[1], 10)
+        self.A_nnz = 32   
+        self.A_idx = gencoordinates(*self.A_shape, self.A_nnz, device='cuda')
+        self.A_val = torch.randn(self.A_nnz, dtype=torch.float64, device='cuda')
+        self.As_coo = torch.sparse_coo_tensor(self.A_idx, self.A_val, self.A_shape, requires_grad=True).coalesce()
+        self.As_csr = self.As_coo.to_sparse_csr()
+        self.Ad = self.As_coo.to_dense()
+        
+        self.Bd = torch.randn(*self.B_shape, dtype=torch.float64, requires_grad=True, device='cuda')
+        self.matmul = sparse_mm
+        
+    def test_matmul_forward_coo(self):
+        x = self.matmul(self.As_coo, self.Bd)
+        self.assertIsInstance(x, torch.Tensor)
+        
+    def test_matmul_forward_csr(self):
+        x = self.matmul(self.As_csr, self.Bd)
+        self.assertIsInstance(x, torch.Tensor)
+        
+        
+    def test_matmul_gradient_coo(self):
+        # Sparse matmul:
+        As1 = self.As_coo.detach().clone()
+        As1.requires_grad = True
+        Bd1 = self.Bd.detach().clone()
+        Bd1.requires_grad = True
+        As1.retain_grad()
+        Bd1.retain_grad()
+        x = self.matmul(As1, Bd1)
+        loss = x.sum()
+        loss.backward()
+        
+        # torch dense matmul:
+        Ad2 = self.Ad.detach().clone()
+        Ad2.requires_grad = True
+        Bd2 = self.Bd.detach().clone()
+        Bd2.requires_grad = True
+        Ad2.retain_grad()
+        Bd2.retain_grad()
+        x_torch = Ad2 @ Bd2
+        loss_torch = x_torch.sum()
+        loss_torch.backward()
+        
+        nz_mask = As1.grad.to_dense() != 0.0
+        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask]).all())
+        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad).all())
+                
+    def test_matmul_gradient_csr(self):
+        # Sparse solver:
+        As1 = self.As_csr.detach().clone()
+        As1.requires_grad = True
+        Bd1 = self.Bd.detach().clone()
+        Bd1.requires_grad = True
+        As1.retain_grad()
+        Bd1.retain_grad()
+        x = self.matmul(As1, Bd1)
+        loss = x.sum()
+        loss.backward()
+        
+        # torch dense solver:
+        Ad2 = self.Ad.detach().clone()
+        Ad2.requires_grad = True
+        Bd2 = self.Bd.detach().clone()
+        Bd2.requires_grad = True
+        Ad2.retain_grad()
+        Bd2.retain_grad()
+        x_torch = Ad2 @ Bd2
+        loss_torch = x_torch.sum()
+        loss_torch.backward()
+        nz_mask = As1.grad.to_dense() != 0.0
+        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask]).all())
+        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad).all())
+        
+if __name__ == "__main__":
+    unittest.main()
\ No newline at end of file
diff --git a/tests/test_sparse_solve.py b/tests/test_sparse_solve.py
new file mode 100644
index 0000000..9d9f755
--- /dev/null
+++ b/tests/test_sparse_solve.py
@@ -0,0 +1,176 @@
+import torch
+import unittest
+from random import randrange
+from torchsparsegradutils import sparse_triangular_solve
+
+def gencoordinates_square_tri(n, ni, upper=True, device='cuda'):
+    """Used to genererate ni random unique off-diagonal coordinates 
+    for upper or lower triangular sparse matrix with size [n, n] """
+    coordinates = set()
+    while True:
+        r, c = randrange(n), randrange(n)
+        if (r < c and upper) or (r > c and not upper):
+            coordinates.add((r, c))
+        if len(coordinates) == ni:
+            return torch.stack([torch.tensor(co) for co in coordinates], dim=-1).to(device)
+        
+def gencoordinates_square_diag(n, device='cuda'):
+    """Generate diagonal indices for square matrix with size [n, n]"""
+    d_idx = torch.arange(n, device=device)
+    d_idx = torch.stack([d_idx, d_idx], dim=0).to(device)
+    return d_idx
+
+class SparseLinearSolveTest(unittest.TestCase):
+    """Test Triangular Linear Sparse Solver for COO and CSR formatted sparse matrices"""
+    def setUp(self) -> None:
+        self.RTOL = 1e-3
+        self.A_shape = (16, 16)  # square matrix
+        self.B_shape = (self.A_shape[1], 10)
+        self.A_nnz = 32  # excluding diagonal of A which will be + A_size nnz
+        self.A_idx = torch.cat([gencoordinates_square_diag(self.A_shape[0]),
+                                gencoordinates_square_tri(self.A_shape[0], self.A_nnz)], dim=1)
+        self.A_val = torch.cat([torch.rand(self.A_shape[0], dtype=torch.float64, device='cuda'),  # [0, 1) to avoid 0 on diag
+                                torch.randn(self.A_nnz, dtype=torch.float64, device='cuda')])
+        
+        self.As_coo_triu = torch.sparse_coo_tensor(self.A_idx, self.A_val, self.A_shape, requires_grad=True).coalesce()
+        self.As_coo_tril = self.As_coo_triu.t()
+        self.As_csr_triu = self.As_coo_triu.to_sparse_csr()
+        self.As_csr_tril = self.As_coo_tril.to_sparse_csr()
+        
+        self.Ad_triu = self.As_coo_triu.to_dense()
+        self.Ad_tril = self.As_coo_tril.to_dense()
+        
+        self.Bd = torch.randn(16, 4, dtype=torch.float64, device='cuda')
+        self.solve = sparse_triangular_solve
+        
+    def test_solver_result_coo_triu(self):
+        x = self.solve(self.As_coo_triu, self.Bd, upper=True)
+        x2 = torch.linalg.solve_triangular(self.Ad_triu, self.Bd, upper=True)
+        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
+        
+    def test_solver_result_coo_tril(self):
+        x = self.solve(self.As_coo_tril, self.Bd, upper=False)
+        x2 = torch.linalg.solve_triangular(self.Ad_tril, self.Bd, upper=False)
+        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
+        
+    def test_solver_result_csr_triu(self):
+        x = self.solve(self.As_csr_triu, self.Bd, upper=True)
+        x2 = torch.linalg.solve_triangular(self.Ad_triu, self.Bd, upper=True)
+        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
+        
+    def test_solver_result_csr_tril(self):
+        x = self.solve(self.As_csr_tril, self.Bd, upper=False)
+        x2 = torch.linalg.solve_triangular(self.Ad_tril, self.Bd, upper=False)
+        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
+        
+    def test_solver_gradient_coo_triu(self):
+        # Sparse solver:
+        As1 = self.As_coo_triu.detach().clone()
+        As1.requires_grad = True
+        Bd1 = self.Bd.detach().clone()
+        Bd1.requires_grad = True
+        As1.retain_grad()
+        Bd1.retain_grad()
+        x = self.solve(As1, Bd1, upper=True)
+        loss = x.sum()
+        loss.backward()
+        
+        # torch dense solver:
+        Ad2 = self.Ad_triu.detach().clone()
+        Ad2.requires_grad = True
+        Bd2 = self.Bd.detach().clone()
+        Bd2.requires_grad = True
+        Ad2.retain_grad()
+        Bd2.retain_grad()
+        x2 = torch.linalg.solve_triangular(Ad2, Bd2, upper=True)
+        loss_torch = x2.sum()
+        loss_torch.backward()
+        
+        nz_mask = As1.grad.to_dense() != 0.0
+        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
+        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
+        
+    def test_solver_gradient_coo_tril(self):
+        # Sparse solver:
+        As1 = self.As_coo_tril.detach().clone()
+        As1.requires_grad = True
+        Bd1 = self.Bd.detach().clone()
+        Bd1.requires_grad = True
+        As1.retain_grad()
+        Bd1.retain_grad()
+        x = self.solve(As1, Bd1, upper=False)
+        loss = x.sum()
+        loss.backward()
+        
+        # torch dense solver:
+        Ad2 = self.Ad_tril.detach().clone()
+        Ad2.requires_grad = True
+        Bd2 = self.Bd.detach().clone()
+        Bd2.requires_grad = True
+        Ad2.retain_grad()
+        Bd2.retain_grad()
+        x2 = torch.linalg.solve_triangular(Ad2, Bd2, upper=False)
+        loss_torch = x2.sum()
+        loss_torch.backward()
+        
+        nz_mask = As1.grad.to_dense() != 0.0
+        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
+        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
+                
+    def test_solver_gradient_csr_triu(self):
+        # Sparse solver:
+        As1 = self.As_csr_triu.detach().clone()
+        As1.requires_grad = True
+        Bd1 = self.Bd.detach().clone()
+        Bd1.requires_grad = True
+        As1.retain_grad()
+        Bd1.retain_grad()
+        x = self.solve(As1, Bd1, upper=True)
+        loss = x.sum()
+        loss.backward()
+        
+        # torch dense solver:
+        Ad2 = self.Ad_triu.detach().clone()
+        Ad2.requires_grad = True
+        Bd2 = self.Bd.detach().clone()
+        Bd2.requires_grad = True
+        Ad2.retain_grad()
+        Bd2.retain_grad()
+        x2 = torch.linalg.solve_triangular(Ad2, Bd2, upper=True)
+        loss_torch = x2.sum()
+        loss_torch.backward()
+        
+        nz_mask = As1.grad.to_dense() != 0.0
+        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
+        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
+        
+    def test_solver_gradient_csr_tril(self):
+        # Sparse solver:
+        As1 = self.As_csr_tril.detach().clone()
+        As1.requires_grad = True
+        Bd1 = self.Bd.detach().clone()
+        Bd1.requires_grad = True
+        As1.retain_grad()
+        Bd1.retain_grad()
+        x = self.solve(As1, Bd1, upper=False)
+        loss = x.sum()
+        loss.backward()
+        
+        # torch dense solver:
+        Ad2 = self.Ad_tril.detach().clone()
+        Ad2.requires_grad = True
+        Bd2 = self.Bd.detach().clone()
+        Bd2.requires_grad = True
+        Ad2.retain_grad()
+        Bd2.retain_grad()
+        x2 = torch.linalg.solve_triangular(Ad2, Bd2, upper=False)
+        loss_torch = x2.sum()
+        loss_torch.backward()
+        
+        nz_mask = As1.grad.to_dense() != 0.0
+        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
+        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
+
+    
+if __name__ == "__main__":
+    unittest.main()
\ No newline at end of file
diff --git a/tests/test_util.py b/tests/test_util.py
deleted file mode 100644
index 3540545..0000000
--- a/tests/test_util.py
+++ /dev/null
@@ -1,16 +0,0 @@
-import unittest
-
-# import torch
-
-import torchsparsegradutils
-
-# # set deterministic seed
-# torch.manual_seed(15)
-
-
-class TestDummyInterface(unittest.TestCase):
-    def test_dummy_shape(self):
-        torchsparsegradutils.dummy_interface()
-
-if __name__ == "__main__":
-    unittest.main()
\ No newline at end of file