Access to previous optimizer state

[kovnerd]

I am trying to add a termination condition to my optimizer loop based on how much the loss value has changed since the last iteration. I am not sure how I would do such a thing. Using the lbfgs tutorial as an example, I trying to implement a "relative function reduction" condition, where given an externally supplied number "factr", the current loss value "f_k", the previous loss value "f_{k-1}", the loop ends if |f_{k}-f_{k-1}| / max(|f_{j}|,|f_{k-1}|,1) <= factr*.

[rdyro]

In optax you typically manage the optimization loop yourself, so you can keep track of the previous loss value as e.g., l_prev and break out of a loop based on that, e.g., using a Python optimization loop or jax.lax.while_loop.

Did you have something else in mind?