Contributing a differentiable and jittable sparse linear solver to JAX core?

[arpastrana]

Hello all! @denizokt and I have worked recently to implement a direct sparse linear solver that is jittable and differentiable via a custom VJP that leverages the implicit function theorem. The solver works on both CPU and GPU. You can check it out here: https://github.com/arpastrana/jax_fdm/blob/main/src/jax_fdm/equilibrium/sparse.py

We put the solver together in the context of a research project on differentiable mechanics for inverse design. As often happens, many mechanical simulations rely on sparse matrices to calculate the structural response of a system efficiently. We identified that there was limited support for direct differentiable sparse solvers (especially for CPU) in JAX at the time, so we decided to take a stab at it.

We think this solver could be useful to other JAX + mechanics folks like us! Are there any chances to contribute our differentiable sparse solver to JAX core? Happy to hear your thoughts and suggestions.

Thanks!

[froystig]

Have you come across JAXopt (https://github.com/google/jaxopt) or lineax? These might be the best places to look for overlap and discuss. You can also read our recent thinking on JAX numpy/scipy scope here.

cc @mblondel @fabianp @patrick-kidger @jakevdp

[arpastrana]

Thanks for the feedback, @froystig!

Lineax looks like a good fit since that library specializes on solving systems of the form Ax=b --the same focus that our direct sparse solver has.

What do you think, @patrick-kidger? Would there be room for us to contribute our work there?

[patrick-kidger]

What kind of system is this solving? (Looking at the implementation it looks like some kind of block sparse system, but I'm not sure exactly what?)

[arpastrana]

The function sparse_solve solves a linear system of the type Ax=b, where A is a sparse matrix and the right-hand side b is a vector or a matrix. This solver is jittable and (reverse-mode) differentiable on CPU and GPU.

Our main contribution here has been to derive the backward pass of the sparse linear solve via the adjoint method (sparse_solve_bwd), and express it in terms of a call to the forward pass (sparse_solve_fwd, the function that solves Ax=b).
On CPU, we use scipy.sparse.linalg.spsolve as the forward pass function (we wrap it in a jax.pure_callback to make it compatible with jit). On GPU, we call jax.experimental.sparse.linalg.spsolve. Both forward and backward passes are registered in a jax.custom_vjp.

As you saw in our code, we assemble a block sparse matrix for A and ravel b in the event that b is a matrix and the linear solve is executed on GPU. This is because the GPU solver implementation in jax.experimental.sparse.linalg.spsolve relies on cuSOLVER, which only offers support for 1D arrays on the right-hand side (RHS) of the linear system.

This is not the case with the CPU solver implementation, which is able to solve the sparse linear system, even in the presence of a RHS matrices, with a single factorization.

[patrick-kidger]

Ah, gotcha. So for the backward pass, Lineax actually already synthesises that for any kind of solver. (The transposition is treated as a property of the matrix instead.)

I think you could try writing this as a Lineax solver (just subclass lx.AbstractLinearSolver like any of the other existing solvers), and you should find that the JVPs+VJPs work automatically. FWIW we haven't yet introduced sparse solvers since JAX's support for sparse operations isn't itself very developed; it might need some more thought how best to handle these cleanly.

[DoTulip]

Hello @arpastrana

I am keenly interested in trying out the sparse solver library you have developed. May I ask what are the requirements for the format of matrix A? Would it be possible to use jax.sparse.BCOO when defining A? It would be highly appreciated if you could provide a simple example to illustrate its usage.

Thank you in advance for your help.