Add example using preconditioning.

[patrick-kidger]

No description provided.

[frederic-bender]

Hey, I'm using a Newton-GMRES-algorithm for implicitly solving large ODEs, based on Diffrax/Optimistix/Lineax – a pretty awesome set of libraries!

Right now I'm trying to figure out, how to create a Jacobi-preconditioner (or a more advanced one) for matrix-free computations. Using lx.diagonal(...) materializes the Jacobian and therefore is not possible for a vector of length ~500,000. Do you have any idea how to solve this problem?

[patrick-kidger]

This depends on the kind of operator you have. A tridiagonal linear operator, for example, can get the diagonal without further materialization!

However it's a classic (sad) fact in autodifferentiation that accessing the diagonal of a Jacobian requires computing the whole Jacobian.

You'll probably need to find a different approach for your problem.

Sorry I couldn't be more help!

[vboussange]

Hey @patrick-kidger, awesome package! It would be indeed very nice to have an example using preconditioning.

Is there any project such as AlgebraicMultigrid.jl or pyamg running in JAX?

Otherwise, I could suggest an example using pyamg.multigrid for solving a Poisson 2D system, see e.g. here.

[patrick-kidger]

I don't thnk so! I'd be very happy to take a PR adding an example of this nature :)

[vboussange]

Ok. Should this go into docs/examples/preconditioning.ipynb or similar? Are we happy to have a dependency on pyamg?

[patrick-kidger]

Yup, exactly there.

I think it'd be best to avoid a dependency on pyamg if at all possible.