sparse jacobian solve

Can the implicit solvers in diffrax use a sparse matrix solve for the jacobian? I'm putting together a benchmark with a few different ode solvers, including diffrax, and the problem in question is stiff (its based on the robertson ODE problem from your examples) and has a small block-diagonal jacobian structure (each block is 3x3 and the total matrix size goes up to 10,000, so very sparse). Diffrax is slowing down significantly at larger matrix sizes and my assumption is that you are using a dense linear solver? Is there a way of swapping this out for a sparse solver? My understanding is that sparse matrix support is still rather experimental for JAX.

You can see the benchmark and results here: https://martinjrobins.github.io/diffsol/benchmarks/python.html. Please feel free to let me know if I'm not using diffrax correctly. I've not used it in anger before, so its entirely possible I'm doing something stupid.

Start by reproducing the stiff Robertson ODE benchmark linked in the issue, focusing on the implicit solvers and their Jacobian behavior at larger matrix sizes. Then investigate whether Diffrax can use a sparse matrix solve, considering the issue's note that sparse matrix support in JAX may be experimental. Done means establishing whether a supported sparse-solver path exists or what design work would be required.