lineax linear operators

[quattro]

Hi 👋 ,

The recent lineax project seems to support a wider range of linear operators and linear solvers that are largely compatible with the approaches outlined internally in jaxopt. I'm curious if there would be an interest or effort in porting internal code to make use of the lineax routines, rather that the custom-rolled solutions for jaxopt.

This may open up the doors for more flexible solvers downstream for root finding.

[vroulet]

Hello @quattro !

Sincere apologies for the late reply. The lineax package definitely has some great features. For example, the implementation of gmres in half-precision with JAX appeared to be unstable and having a better implementation like the one in lineax would be great. For now, the philosophy of JAXopt was to try avoiding additional packages and be dependent only on JAX, I believe. For example, while we have an interface to use Optax, we don't have it as a direct dependency. So having additional linear operators from another package was not sought after. However, as said above, it would be great to update some linear solvers and integrate some efficiency brought up by lineax. There should be a discussion in September about such merges.

Would it be possible to discuss such integration with the JAX team too? I believe revamping some of the jax.sparse.linalg package with some of the features of lineax would be great too and would avoid fragmenting the JAX ecosystem.

Again sorry for the delay, your proposal would lead to substantial changes and all of us are not available right now to discuss it. We'll reboot this conversation in September at the latest.

Thanks for initiating this!