Update for latest ParameterSchedulers.jl release

[darsnack]

Due to some finagling with #1506 and work on Optimisers.jl, I made some changes to the names/API of ParameterSchedulers.jl. This just updates the docs to reflect those names.

PR Checklist

• Tests are added
• Entry in NEWS.md
• Documentation, if applicable
• Final review from @dhairyagandhi96 (for API changes).

[CarloLucibello]

The docs are fine, maybe just briefly mention what Stateful does since it is not very clear.

Unrelated to this PR: do we really need a wrapper around a scheduler to make it stateful? Cannot they just contain the counter, then allow either stateless or stateful behavior depending on the context?

[darsnack]

We could do it that way, but I think keeping the internal state in sync will get annoying/error-prone. Plus, I think Optimisers.jl will want everything to be immutable.

Instead, I suggest that we do something like:

Exp(args...; stateful=false) = Stateful(Exp(args...))

Default can be whatever we decide is most convenient.

[CarloLucibello]

I'm growing to dislike more and more this whole functional/stateless story, makes everything more annoying and verbose without providing any clear advantage. Anyways, the convenience constructor seems fine. I will vote for stateful=true and likely be outvoted 😞

[CarloLucibello]

do you want to update ParameterSchedulers first or merge this?

[ToucheSir]

Not to derail this too much, but I wouldn't care if stateful=true or stateful=false iff. Stateful remains non-iterable**. Stateful iterators (i.e. those that self-mutate on iteration) are a big footgun, and forcing users to explicitly advance via next! would all but remove that footgun.

** if you look at code using JAX libraries and Optax, they're not shy about mutating hyperparams in the training loop either. Really the immutability requirement is more for anything that might have to run on an accelerator (and thus needs to be XLA friendly)

[darsnack]

FWIW, I didn't make the schedules immutable because of the whole Optimisers.jl story (if it were up to me we'd still have apply!...). I basically looked at a schedule and decided that the most intuitive way to think of one in Julia is as an iterator:

s = [lambda * gamma^(t - 1) for t in 1:total_iters]

Then I tried to get something that behaved that way but without being so verbose. Since iterators in Julia are explicit state, this basically meant schedules would be explicit state.

I'm biased because I don't use Stateful. In my experience, you either schedule every epoch or every mini-batch. For the former, I use zip, and the latter I use Scheduler. But I know some people would prefer next! which is why I wrote Stateful. I'm happy for the default to be stateful=true as well.

[darsnack]

do you want to update ParameterSchedulers first or merge this?

Let me do a patch release just to make sure the proposed syntax works. Then I'll update this to reflect it.

[darsnack]

@CarloLucibello I ended up just rephrasing the docs in the PR as originally planned. I wasn't happy with stateful=true overshadowing the default constructors (thus, overshadowing most of the functionality in the package). I could do stateful=false, but I figured we can discuss a better interface in #1506. I didn't want to submit a bunch of ad-hoc patch releases as we iterate.

[darsnack]

Thanks!

bors r+

[bors]

Build succeeded:

• buildkite/flux-dot-jl