Using sksparse.cholmod.cholesky for solving sparse systems

username: Elling_Svee
cooked: <p>Hi!</p>
<p>I use <code>torch.optim</code> to optimize a function which actively uses sparse matrix operations. The current <code>torch.sparse</code> package does not support operations like computing Cholesky factorizations of sparse tensors, computing log-determinants or solving sparse linear systems. Therefore it would be very beneficial if I could find some way of implementing these operations as custom autograd functions.</p>
<p>From previous experience, I have found that <code>sksparse.cholmod.cholesky</code> works very well for sparse Cholesky decomposition of <code>scipy.sparse</code> matrices. This also supports computations of the log-determinant and can be used to solve systems. I therefore want to use this within my autograd function.</p>
<p>Following this tutorial (<a href="https://blog.flaport.net/solving-sparse-linear-systems-in-pytorch.html" class="inline-onebox" rel="noopener nofollow ugc">Flaport.net | Creating a Pytorch solver for sparse linear systems</a>) I have the following code that solves a sparse system using the <code>scipy.sparse.linalg.spsolve</code> solver.</p>
<pre><code class="lang-auto">import torch
import numpy as np
import scipy.sparse as sparse

class Solve(torch.autograd.Function):
    @staticmethod
    def forward(ctx, A, b):
        A_np = A.data.numpy()
        b_np = b.data.numpy()
        A_sp = sparse.csc_matrix(A_np)
        
        # Solve using spsplve
        x_np = sparse.linalg.spsolve(A_sp, b_np)

        x = torch.tensor(x_np, requires_grad=True)
        ctx.save_for_backward(A, b, x)
        return x

    @staticmethod
    def backward(ctx, grad):
        A, b, x = ctx.saved_tensors
        gradb = Solve.apply(A.T, grad)
        gradA = -gradb[:, None] * x[None, :]
        return gradA, gradb

custom_solve_system = Solve.apply

# Test the solver
A = torch.randn(3, 3, requires_grad=True, dtype=torch.float64)
A = (A + A.t()) / 2 + torch.eye(3) # Make SPD
b = torch.randn(3, requires_grad=True, dtype=torch.float64)

# Returns True
torch.autograd.gradcheck(custom_solve_system, [A.to_dense(), b])
</code></pre>
<p>According to the <code>gradcheck</code> function, the <code>custom_solve_system</code> function works as intended. However, now I replace the <code>spsolve</code> with a solver using the <code>sksparse.cholmod.cholesky</code>. This apparently causes problems.</p>
<pre><code class="lang-auto">import torch
import numpy as np
import scipy.sparse as sparse
from sksparse.cholmod import cholesky

class Solve2(torch.autograd.Function):
    @staticmethod
    def forward(ctx, A, b):
        A_np = A.data.numpy()
        b_np = b.data.numpy()
        A_sp = sparse.csc_matrix(A_np)

        # Instead use the sksparse Cholesky factorization
        A_chol = cholesky(A_sp)
        x_np = A_chol.solve_A(b_np)

        x = torch.tensor(x_np, requires_grad=True)
        ctx.save_for_backward(A, b, x)
        return x

    @staticmethod
    def backward(ctx, grad):
        A, b, x = ctx.saved_tensors
        gradb = Solve2.apply(A.T, grad)
        gradA = -gradb[:, None] * x[None, :]
        return gradA, gradb

custom_solve_system_2 = Solve2.apply

# Test the solver
A = torch.randn(3, 3, requires_grad=True, dtype=torch.float64)
A = (A + A.t()) / 2 + torch.eye(3) # Make SPD
b = torch.randn(3, requires_grad=True, dtype=torch.float64)

# Raises an error.
torch.autograd.gradcheck(custom_solve_system_2, [A.to_dense(), b])
</code></pre>
<p>Since this gives an  error, the function no longer works according to the <code>gradcheck</code>. Does anybody know how I can fix this? Any help would be much appreciated!</p>