Linear solver benchmarking on EqQPs about jaxopt HOT 5 OPEN

Interesting results! I would be also interested in a benchmark of JAX gmres vs. SciPy gmres.

from jaxopt.

I don't really have anything useful to add, but I found that cg and bicg were superior wrt to stability & much better wrt runtime on large sparse and large sparse "slightly singular" systems- quadratics involving a graph Laplacian.

I didn't have much success preconditioning any of them.

from jaxopt.

Yeah, the reason I'm not using CG is because the EqQP is indefinite -- so I'd have to use CG on the normal system A.T A = A.T b, which can have a really big condition number.

from jaxopt.

The restart parameter for GMRES is critical. Here's a new plot with restart=625 for this example.

LGMRES is an automated way of finding the restart parameter, amongst other things).

Perf for GMRES is now on par with LGMRES.

from jaxopt.

@mblondel sorry I didn't get around to that benchmark. I recently found this JAX issue which says that GMRES is much slower on GPU than on CPU for Jax.
I'm now wondering how to wrap the GMRES to move things to CPU, including for the backward pass, if it's used as a submodule in a longer differentiable pipeline.

from jaxopt.