Update termination logic

[bagibence]

Issue:

Please correct me if this is intentional, but I noticed that the termination/continuation in _iterate might have a bug.
I ran into an issue where non-finite values in the initial parameters lead to an optimization running until max_steps.

Currently solvers initialize result=RESULTS.successful and change it if there is a failure of some kind (e.g. parameter divergence).
I would assume that in those failure cases optimization would stop. However, with continuing if jnp.invert(terminate) | (result != RESULTS.successful), it keeps running until max_steps is reached.

Minimal example reproducing the problem:

import jax.numpy as jnp
import optimistix as optx

def loss(params, args):
    return optx.two_norm(params)

init_params = jnp.array(jnp.inf) # runs for 1000 iterations on dev
# init_params = jnp.array(5.0) # runs fine

solver = optx.BFGS(1e-8, 1e-8, optx.two_norm)
sol = optx.minimise(loss, solver, init_params, args=(), max_steps=1000, throw=False)
print(sol.stats["num_steps"])
print(sol.result)
print(sol.value)

Interestingly, with JAXopt I didn't run into this because their continuation criterion was false if state.error was NaN -- given by something like two_norm(inf - inf) in these cases.

Changes:

• Updated _iterate to continue if jnp.invert(terminate | (result != RESULTS.successful)), so stop if the solver says so or not successful.
• Updated some solvers to terminate and update the result if the parameters are non-finite. I haven't gone through all solvers. Please let me know if others need to be updated or if this should be included in _iterate instead (see todo).
• Added a test for stopping on divergence. Although now I realize that the example above might be simpler. Let me know if you'd like to use that instead.
• Updated newton_chord to only stop after 2 iterations are done and its test.

Additional comment:

While I think the current solution works, for me it is unintuitive that solvers have a terminate method, but that does not actually fully determine if they stop or not -- as _iterate looks at both terminate and results returned by the solver.
My ideal solution would be that _iterate stops only based on the value of terminate and the solver has all control over that. But I think that would require updating each solver's terminate method, so I wanted to run this by you before doing that in case I'm wrong here.

[patrick-kidger]

Thanks for opening this! I wonder if this should be a thing we check every step, or just at initialisation? I'm guessing that checking every step may be moderately expensive. If we just do it at initialisation then I think that would also be something that we can just do for every solver. WDYT?

While I think the current solution works, for me it is unintuitive that solvers have a terminate method, but that does not actually fully determine if they stop or not -- as _iterate looks at both terminate and results returned by the solver.

On this topic, this choice is essentially just to help out authors of custom solvers. A non-successful result is always intended to indicate a failure case, and this way it's not possible to forget to check it.

[bagibence]

I would say it's definitely needed on initialization, but also on every (n-th) iteration to catch divergence and stop early. Not sure about how expensive the check is, it should be linear in the number of parameters.
It could also be done at the end as a diagnostic to avoid falsely reporting success.

A cheaper but less explicit alternative could be hiding the divergence check in cauchy_termination, checking just two values instead of every parameter:

y_diff_norm = norm((ω(y_diff).call(jnp.abs) / y_scale**ω).ω)
f_diff_norm = norm((ω(f_diff).call(jnp.abs) / f_scale**ω).ω)
diverged = jnp.invert(jnp.isfinite(y_diff_norm) & jnp.isfinite(f_diff_norm))

and returning (converged, diverged) or (y_converged, f_converged, y_diverged, f_diverged) for more detail.

A downside of this could be that if y diverged, searches might still iterate forever. Adding a max number of search steps or the same check inside the searches could help with that.

Update: added a link to an implementation of cauchy_termination returning a tuple.