Repository: google/jax

Issue:
[sparse] Feature request: add grad support to experimental.sparse.linalg.spsolve
`experimental.sparse.linalg.spsolve` is implemented via the [`cuda_csrlsvqr` routine](https://github.com/google/jax/blob/7708cf5c21e3296d0a607125f8686d5a9712c1ca/jaxlib/gpu_solver.py#L258). If I am not mistaken, the gradient should be computable (GPU only) with another call to `spsolve`. 


Discussion:
FWIW you should be able to do this by wrapping into a `lineax.AbstractSolver`. ([Lineax is here](https://github.com/google/lineax).) This automatically synthesises the backward pass given just the forward solve.
Out of curiosity, how do you automatically construct the transpose rule for an abstract solver if no transpose solver is available?
Transposition is handled as a property of the operator, rather than the solver. I.e. we basically just do `linear_solve(operator.transpose(), solver, ...)`.

So yeah, you still have to define transposition somewhere! But this was (a) we can do single-dispatch on the operator type (e.g. efficiently transpose a tridiagonal matrix), and (b) it doesn't come up when adding new solvers.
Makes sense, thanks! Since `spsolve_p` does not provide any built-in transpose solve, it sounds like a simple `lineax` wrapper would not be sufficient to make it work with a backward pass.

@Alessio-Amaolo hopefully #16938 will enable the things you want to do, even though it has to resort to using a manual CSR transpose for the backward pass.
> Makes sense, thanks! Since spsolve_p does not provide any built-in transpose solve, it sounds like a simple lineax wrapper would not be sufficient to make it work with a backward pass.

I think it should. It's the *matrix* that is transposed, not the solver.
Agreed, I used "transpose solve" as a shorthand for "transpose the matrix and then solve", sorry that wasn't clear.

The issue is that this primitive takes a CSR representation as the input, and unlike COO there's no trivial way to transpose a CSR representation without calling out to some other kernel. Some solvers have a transpose flag built-in to handle this case, but `spsolve_p` does not, because `csrlsvqr` doesn't provide any such flag.
Right, I see! I hadn't fully appreciated that (a) `spsolve` only accepts CSR, and that (b) the transpose rule for a `CSR` matrix wasn't already defined (due to it being so inefficient).

Resolving patch:
diff --git a/jax/experimental/sparse/linalg.py b/jax/experimental/sparse/linalg.py
--- a/jax/experimental/sparse/linalg.py
+++ b/jax/experimental/sparse/linalg.py
@@ -20,10 +20,12 @@
 import jax
 import jax.numpy as jnp
 
+from jax.experimental import sparse
 from jax.interpreters import mlir
 from jax.interpreters import xla
 
 from jax._src import core
+from jax._src.interpreters import ad
 from jax._src.lib import gpu_solver
 
 import numpy as np
@@ -549,9 +551,48 @@ def _callback(data, indices, indptr, b, **kwargs):
   return result
 
 
+def _spsolve_jvp_lhs(data_dot, data, indices, indptr, b, **kwds):
+    # d/dM M^-1 b = M^-1 M_dot M^-1 b
+    p = spsolve(data, indices, indptr, b, **kwds)
+    q = sparse.csr_matvec_p.bind(data_dot, indices, indptr, p,
+                                 shape=(indptr.size - 1, len(b)),
+                                 transpose=False)
+    return -spsolve(data, indices, indptr, q, **kwds)
+
+
+def _spsolve_jvp_rhs(b_dot, data, indices, indptr, b, **kwds):
+    # d/db M^-1 b = M^-1 b_dot
+    return spsolve(data, indices, indptr, b_dot, **kwds)
+
+
+def _csr_transpose(data, indices, indptr):
+  # Transpose of a square CSR matrix
+  m = indptr.size - 1
+  row = jnp.cumsum(jnp.zeros_like(indices).at[indptr].add(1)) - 1
+  row_T, indices_T, data_T = jax.lax.sort((indices, row, data), num_keys=2)
+  indptr_T = jnp.zeros_like(indptr).at[1:].set(
+      jnp.cumsum(jnp.bincount(row_T, length=m)).astype(indptr.dtype))
+  return data_T, indices_T, indptr_T
+
+
+def _spsolve_transpose(ct, data, indices, indptr, b, **kwds):
+  assert not ad.is_undefined_primal(indices)
+  assert not ad.is_undefined_primal(indptr)
+  if ad.is_undefined_primal(b):
+    # TODO(jakevdp): can we do this without an explicit transpose?
+    data_T, indices_T, indptr_T = _csr_transpose(data, indices, indptr)
+    ct_out = spsolve(data_T, indices_T, indptr_T, ct, **kwds)
+    return data, indices, indptr, ct_out
+  else:
+    # Should never reach here, because JVP is linear wrt data.
+    raise NotImplementedError("spsolve transpose with respect to data")
+
+
 spsolve_p = core.Primitive('spsolve')
 spsolve_p.def_impl(functools.partial(xla.apply_primitive, spsolve_p))
 spsolve_p.def_abstract_eval(_spsolve_abstract_eval)
+ad.defjvp(spsolve_p, _spsolve_jvp_lhs, None, None, _spsolve_jvp_rhs)
+ad.primitive_transposes[spsolve_p] = _spsolve_transpose
 mlir.register_lowering(spsolve_p, _spsolve_gpu_lowering, platform='cuda')
 mlir.register_lowering(spsolve_p, _spsolve_cpu_lowering, platform='cpu')
 


Test patch:
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -2799,6 +2799,28 @@ def sparse_solve(data, indices, indptr, b):
     self.assertAllClose(a @ x, b, rtol=1e-2, atol=1e-3)
     self._CompileAndCheck(sparse_solve, args_maker)
 
+  @jtu.sample_product(
+    size=[10, 20, 50],
+    dtype=jtu.dtypes.floating,
+  )
+  @unittest.skipIf(jtu.device_under_test() == "tpu", "test requires CPU or GPU")
+  @unittest.skipIf(jtu.device_under_test() == "cuda" and not GPU_LOWERING_ENABLED,
+                   "test requires cusparse/cusolver")
+  @jtu.skip_on_devices("rocm", "test requires cusolver")
+  def test_sparse_qr_linear_solver_grads(self, size, dtype):
+    rng = rand_sparse(self.rng())
+    a = rng((size, size), dtype)
+    nse = (a != 0).sum()
+    data, indices, indptr = sparse_csr._csr_fromdense(a, nse=nse)
+
+    rng_k = jtu.rand_default(self.rng())
+    b = rng_k([size], dtype)
+
+    def sparse_solve(data, b, tol=1e-8):
+      return sparse.linalg.spsolve(data, indices, indptr, b, tol=tol)
+
+    jtu.check_grads(sparse_solve, (data, b), order=1, rtol=0.05, atol=0.05)
+
 
 class SparseUtilTest(sptu.SparseTestCase):