Sparse matrix-vector products for solvers

[SNMS95]

I am writing a finite element analysis (FEA) code in JAX. In order to solve the system of equations (KU=F), I currently use the Jacobian of the residual to find the 'K' (stiffness matrix). This is easy in JAX since the residual is a differentiable function.

To solve the system of equations, I currently use the jvp function on the residual to get the matvec to use in the forward solver. This is as follows:

def jvp_mat_vec(v):
  _, Kv = jax.jvp(residual_fn, primals, v)
return Kv

u_solution = solver(jvp_fn, b)  # Some solver

But someone told me that since K is sparse, I should not use the JVP but instead create a sparse matrix out of K and create an explicit matrix-vector product function. This is shown below

  K_sp = BCOO.from_scipy_sparse(K_sp_scipy).sort_indices()
  def compute_linearized_residual(u):
      return K_sp @ u

u_solution = solver(compute_linearized_residual, b) # Some solver

Is the second one really faster ? This is important for me since the matvec function will be called many times in the solver.

[jakevdp]

In general sparse operations will not be faster than dense operations, particularly on accelerators like GPU (this is not just true in JAX, but in virtually all systems that run on modern hardware). But sparse operations can be useful if the dense version of your matrix is too large to fit in memory, or if your matrix is extremely sparse (i.e. ~99.9% sparse) so that the indexing overhead does not dominate the cost of a dense matmul.

[SNMS95]

Oh.. So, I can use the jax.experimental.sparse matrices for saving GPU memory but not the actual speed of computation.

Thanks a lot jake!

[patrick-kidger]

To go into a little more detail: there's several facets to which of these is best. For example, if your residual_fn is very long, then it may simply take a long time to run through all of its operations in the jvp version.

Another point to consider: if you call jvp_mat_vec in multiple places, then you may find that your compile time explodes as each call has to be compiled separately.

Here's a third point: with jax.jvp, you are recomputing the primal every time, just to throw it away. Instead, you may wish to trade memory for computational work, by using jax.linearize to compute the primal just once. This will then save all of the residuals, and you will only need to evaluate the tangent sweep when you evaluate the linearised function.

There's all kinds of trade-offs to be made here, so generally speaking you'll need to try both (or know something about the structure of your problem) to figure out which is best wrt runtime / compile time / memory usage.

Finally, allow me to point you at Lineax, which is our enhanced suite of linear solvers for JAX. In particular it includes a number of kinds of linear operator, which would allow you to switch between the jvp / linearize / sparse approaches straightforwardly.