pre_update for OptaxSolver run function? about jaxopt HOT 6 CLOSED

I'm doing some changes in the base classes, I will tackle this issue too.

from jaxopt.

Indeed, pre_update is currently only used in run_iterator and not in run. I was a bit hesitant about this, I am not sure what are the right semantics here. If we add it to run, then all solvers, including non-stochastic ones, will need a pre_update option. Is it worth it for batch solvers?

Another option would be to apply pre_update inside update. This means that every solver with a pre_update option would be in charge of calling it. This way, both run and run_iterator would use pre_update as a result. And this way pre_update doesn't 'leak' everywhere in the code base.

So the question is: do we make pre_update an integral part of our solver API or not?

Any opinion on this would be welcome.

from jaxopt.

I may be misinterpreting the role of pre_update (in particular, why it is significantly more useful in stochastic solvers).
My main use of the pre_update was tracking the progress of some running stats as well as how far along it is within the optimization (to know how much longer I have to wait). If pre_update is the only way to do this, I think it should be added everywhere, since tracking stats&progress seems pretty useful every time we have a solver. If we don't add it then I would put some sort of warning (or change the name?) in the OptaxSolver definition, because intuitively one would expect the pre_update to run if you give it to the OptaxSolver, whenever you call update.

Happy to chat more about it!

from jaxopt.

Moving the call to pre_update to within OptaxSolver.update sounds good. Do you want to do it?

from jaxopt.

Done in #29, opinions welcome!

from jaxopt.

Awesome! Thanks for solving it :)

from jaxopt.