google / google/jaxopt

Linear solver benchmarking on EqQPs

Open
#168 5 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.1k
Forks
76
Avg merge
2d 21h
Merged PRs (30d)
1

Description

I benchmarked GMRES (currently default in jax.eq_qp) vs other `scipy` solvers: [minres](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.minres.html#scipy.sparse.linalg.minres) and [LGMRES](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.lgmres.html#scipy.sparse.linalg.lgmres). I sampled random equality constrained QP KKT matrices and targets, and found pretty stark differences between the solvers. Primal dim: 1500, dual dim: 1000.

TLDR: LGMRES is much faster, and GMRES is the slowest of the three. NB: Minres requires the matrix in the linear system to be symmetric (although indefinite), while gmres and lgmres don't.

I recall a discussion about linear solvers living in the `jaxopt` package vs `jax.scipy.sparse.linalg`.
I don't have time right now to implement LGMRES in jax, (and would be grateful if anyone has the bandwidth and the interest) but I will try to get around to it in the next few weeks. Should the solver's code live in jaxopt?

Colab for repro: https://colab.research.google.com/drive/1Ge1-gmuknDQq0rHpnSrHpvwLbU23i6YG?usp=sharing

![image](https://user-images.githubusercontent.com/11814989/152444223-57ba6bea-0c2c-4eb4-aa62-afa9ab70e575.png)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.