Support schedules for preconditioners
Created by: PhilipVinc
Remake of #1142, with a workaround in order not to throw a warning at every iteration.
This could later be slightly extended to also support diag_scale
from the other @attila-i-szabo 's PR. This forces preconditions to accept a third argument, but they can ignore it and it can be not passed around so that SR and everything else can still be used in custom loops without breaking user code.
I also plan in the future to migrate our preconditions to a stateful interface like optax optimisers, but for now I'd like to make minimal changes...
Only tests are missing. If someone wants to speed up this being merged, push some tests because I have little time right now.