pre_update for OptaxSolver run function?

[FerranAlet]

Hi,
Defining a pre_update function works well in OptaxSolver when using the run_iterator function (like in the MNIST example). However, it does not work with the run function. I checked the implementation and indeed there is no call to the pre_update inside 'run' but there is one inside 'run_iterator'. Is this a small bug or is this by design?

Thanks!

[mblondel]

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.

[FerranAlet]

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!

[mblondel]

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

[mblondel]

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

[mblondel]

Done in #29, opinions welcome!

[FerranAlet]

Awesome! Thanks for solving it :)