AMG support?

[MRiabov]

Hello Lineax,
I'm currently developing a multiphysics, multibody solver - kind of like Project Chrono on Jax.

The absolute best solver type on the market would be algebraic multigrid (AMG), and wow should it go with jax.jit and vmap well. It doesn't have an open-source Jax implementation yet, however.

Would it be possible that you implement AMG in some observable future?

Cheers.

[patrick-kidger]

See also the recent discussion around AMG here: #8

Unfortunately I think it's unlikely we'll have an implementation, available by default, in Lineax in the near future. I'm cautious about adding new solvers as they can easily become difficult to maintain.

That said, it should be essentially straightforward for you to use this with Lineax: subclass lineax.AbstractLinearSolver and implement the relevant abstract methods with your chosen implementation of AMG. This will then be compatible with all of Lineax, and all of its usage in downstream libraries (Optimistix, Diffrax, ...)