patrick-kidger
Descent has to know about solver options, at least in some cases
Example: optx.AbstractGaussNewton defaults to forward-mode autodiff, but the descent for IndirectLevenbergMarquardt has a root-finder in it which defaults to Newton and that one does not default to forward-mode autodiff. (Same thing in Dogleg.)
With the current API, we have no way of synchronising these in the background, since options have to be passed to the root_find that is called in the indirect damped Newton descent. Synchronising this would require a custom descent in which options are an attribute.
Now, I can think about whether forward-mode autodiff would be a good default option in the root-finders - we can certainly support it, everything is in place - but the question is: what is a good way to synchronise options for solvers that have some degree of nesting?
If it's specifically around autodiff then I still hold out hope that JAX will eventually come to support forward-through-custom_vjp, and that we can do away with the huge amounts of complexity that exists because of its lack.
As such when I've ran into this before, I've worked around it by hardcoding special handling in the top-level solver: 'if forward mode and have solver Foo, then set solver.jac = "bwd"', this kind of thing.
Do you have other examples where this might be needed?
Fitting parameters of an ODE using diffrax.ForwardMode. But I only came across it when troubleshooting something, normally I use LevenbergMarquardt for this task.
If this is a sometime-fix on the JAX side then we can table this for now!