Adding a preconditioner example
Hey there, here is an attempt to make an example for the use of a preconditioner, as discussed in #8. I implemented a simple Jacobi preconditioner, which does not seem to provide much advantage though. @patrick-kidger Any idea of a low hanging fruit solution to have a nicer demonstration of the use of preconditioners?
I personally would love to see an example with a hand crafted preconditioner such as the one proposed, and a more involved using e.g.pyamg preconditioning capability, with a possible pure_callback call.
In terms of the implementation, I like this example! It's short and easy to follow.
(Nit: jr.PRNGKey -> jr.key; the former has been replaced with the latter.)
Also tagging @frederic-bender from #8, as I don't have strong feelings about what would make for a stronger example of preconditioning. (I note that the Julia folks also only consider a very simple example.)