Add callback functionality to solvers that use verbose_print.

[StephenJamesCashen]

We currently have a complicated home-brewed set of optimizers and we'd like to switch to using optimistix for all the usual reasons. However we have need to analyze and log fit dynamics captured using methods such as jax.debug.callback as part of our development of internal fitting routines.

In this PR I have added a new class variable verbose_callback to the optimistix classes that utilize _misc.verbose_print. This allows end users to override the verbose_print behaviour to implement arbitrary callbacks with jax.debug.callback such as dumping progress to logs, memory etc...

The justification for overriding verbose_print is that consideration has already been taken as to what data is considered appropriate to pass to a jax callback and an API for customizing what data to pass is already in place.

Classes impacted are OptaxMinimizer AbstractGaussNewton and AbstractQuasiNewton. The default callback remains verbose_print and unless a child class explicitly overrides this class variable their behaviour remains the same.

A pattern for a simple use case is demonstrated in the test_verbose.py unit test.

[patrick-kidger]

Thanks for the PR! This basically looks reasonable to me. My main question – is the full set of triples passed to verbose_print useful for your application, or would just the string that's built suffice? (So that you can redirect somewhere other than stdout.)

The verbose printing is pretty hacked-together right now (c.f. the way it's just a bunch of triples keyed by strings...) so my feeling is that the natural user-facing API here is either going to be more specific (consuming just a string as above) or more general (consuming the entire iterating state).

[StephenJamesCashen]

Thanks for the PR! This basically looks reasonable to me. My main question – is the full set of triples passed to verbose_print useful for your application, or would just the string that's built suffice? (So that you can redirect somewhere other than stdout.)

The verbose printing is pretty hacked-together right now (c.f. the way it's just a bunch of triples keyed by strings...) so my feeling is that the natural user-facing API here is either going to be more specific (consuming just a string as above) or more general (consuming the entire iterating state).

That full set of triples is useful - we want to use that to structure the data in the callback. (Log it to a table for instance or a queue in memory)

I am hesitant to faciliate passing back the full state just because it could potentially contain enough information to incur significant overhead to pass to a callback. The verbose api already present is convenient and it contains useful metrics not directly present in the state such as the sum of squared residuals. It covers our needs in terms of information.

[StephenJamesCashen]

One thing that would be useful is allowing verbose_callback to be passed as an instance variable in __init__ where for instance we want to pass runtime context to the callback via closure. @patrick-kidger How would you feel about that API change?

[patrick-kidger]

Okay! Sorry for the delay. Having thought this over a bit, I think I've settled on a design, and have opened #199 as a quick (untested) demo of what I have in mind. I think this should offer a suitable way for you to customise what is displayed, by allowing you to pass a custom callable. WDYT?

(And this includes customising the callback during __init__.)

[StephenJamesCashen]

Sorry for delay, back from vacation now. #199 looks great, would definitely solve our problem. Closing this PR