Repository: FluxML/Optimisers.jl
Pull request #30 (2022-01-27 17:38:35)

Wrap optimiser and state in a struct
As discussed here https://github.com/FluxML/Optimisers.jl/pull/29#discussion_r793256133, this removes the need to explicitly pass the optimiser itself, by wrapping it up with the state. Then `Optimisers.update` takes 3 arguments by default, not 4.

darsnack: Slight bike shedding comment: what about the framing being `Optimiser(rule, state)`? So, `ADAM` etc. are "update rules" that when combined state initialized on some `x` makes a complete "optimizer."
mcabbott: So you want to reserve the word "state" for what `init` returns, and call ADAM a "rule", and call the combined thing an "optimiser".

That might be better. It is a bit weird to use "state" to mean two things.

The current rule descriptions all say things like "version of the ADAM optimiser". Maybe "optimiser" isn't the perfect word for the combined thing, but what is?
darsnack: > Maybe opt = Optimisers.setup(ADAM(), model) would be nice?

This is currently there in form of `init` + `state`. `init` is the low level function that actually returns the state for a leaf. `state` is the user facing function (equivalent to your `setup`). I do think `setup` makes a bit more sense.

> The current rule descriptions all say things like "version of the ADAM optimiser". Maybe "optimiser" isn't the perfect word for the combined thing, but what is?

I think it makes sense. In the current non-functional understanding, our optimizers are just rules + state.

@ToucheSir and I have always pushed that these functional optimizers are really more like gradient transformation rules. Under this framing, things like `WeightDecay` or `ClipGrad` make sense.
mcabbott: How about "store" for the name of the complete thing, rule + state?
darsnack: Hmm seems like a verb being squeezed into a noun. With `setup`, I suppose no one needs to type this name out. It could just be `RuleAndState` or `RuleState`
ToucheSir: +1 for `RuleState`.
mcabbott: I was thinking took "store" as a noun (larder, storehouse, stockpile). RuleState is also OK. It could also be `Setup`, since that's also a noun, perhaps that's even cleaner?
darsnack: > I was thinking took "store" as a noun (larder, storehouse, stockpile).

I guess this is murky due to colloquialisms 😄.

`Setup` is also okay by me if we want something short.
ToucheSir: This is getting into ultra bikeshed territory, but "`update` with `Setup`" reads more awkwardly to me than "`update` with `RuleState`".
mcabbott: Ok, we think a bit. Most importantly, what single letter variable name does it get? 
ToucheSir: If I may so bold to suggest an multi-letter initialism instead. So `st`/`sp` for `Setup`, or `rs` for `RuleState`/`RuleAndState`.
ToucheSir: Want to give Kyle a chance to walk through this, but otherwise looks good. Do we want to include tweaks to https://github.com/FluxML/Optimisers.jl/blob/master/docs/src/index.md in this PR or a follow-up?

Patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index b06d565..ef7bf7a 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -1,7 +1,12 @@
 name: CI
 on:
-  - push
-  - pull_request
+  pull_request:
+    branches:
+      - master
+  push:
+    branches:
+      - master
+    tags: '*'
 jobs:
   test:
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
@@ -11,6 +16,7 @@ jobs:
       matrix:
         version:
           - '1.6'
+          - '1'
           - 'nightly'
         os:
           - ubuntu-latest
@@ -40,6 +46,14 @@ jobs:
         continue-on-error: ${{ matrix.julia-version == 'nightly' }}
       - uses: julia-actions/julia-runtest@v1
         continue-on-error: ${{ matrix.julia-version == 'nightly' }}
+      - uses: julia-actions/julia-processcoverage@v1
+        if: matrix.version == '1' && matrix.os == 'ubuntu-latest'
+        continue-on-error: ${{ matrix.julia-version == 'nightly' }}
+      - uses: codecov/codecov-action@v2
+        if: matrix.version == '1' && matrix.os == 'ubuntu-latest'
+        with:
+          file: lcov.info
+        continue-on-error: ${{ matrix.julia-version == 'nightly' }}
 
   docs:
     name: Documentation
diff --git a/docs/src/api.md b/docs/src/api.md
index 8ef3a33..96e34bf 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -1,3 +1,6 @@
+
+## Optimisation Rules
+
 ```@docs
 Optimisers.Descent
 Optimisers.Momentum
@@ -15,9 +18,27 @@ Optimisers.ADAMW
 Optimisers.AdaBelief
 ```
 
+In addition to the main course, you may wish to order some of these condiments:
+
 ```@docs
 Optimisers.ClipGrad
 Optimisers.ClipNorm
 Optimisers.WeightDecay
 Optimisers.OptimiserChain
 ```
+
+## Model Interface
+
+```@docs
+Optimisers.setup
+Optimisers.update
+Optimisers.update!
+```
+
+## Rule Definition
+
+```@docs
+Optimisers.apply!
+Optimisers.init
+Optimisers.@..
+```
diff --git a/docs/src/index.md b/docs/src/index.md
index c0d23ac..ce87a32 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -2,21 +2,29 @@
 
 ## Define an Optimiser
 
+A new optimiser must overload two functions, `apply!` and `init`:
+
 ```julia
-# Define a container to hold any optimiser specific parameters (if any)
-struct Descent{T}
+# Define a container to hold any optimiser specific parameters (if any):
+struct DecayDescent{T}
   η::T
 end
 
-# Define an `apply!` rule with which to update the current params
-# using the gradients
-function Optimisers.apply!(o::Descent, state, m, m̄)
-    o.η .* m̄, state
+# Define an `apply!` rule which encodes how the gradients will be used to
+# update the parameters:
+function Optimisers.apply!(o::DecayDescent, state, x, x̄)
+  newx̄ = (o.η / √state) .* x̄
+  nextstate = state + 1
+  return nextstate, newx̄
 end
 
-Optimisers.init(o, x::AbstractArray) = nothing
+# Define the function which sets up the initial state (if any):
+Optimisers.init(o::DecayDescent, x::AbstractArray) = 1
 ```
 
+The parameters will be immediately updated to `x .- newx̄`, while `nextstate` is
+caried to the next iteration.
+
 Notice that the state is handled separately from the optimiser itself. This
 is a key design principle and allows users to manage their own state explicitly.
 
@@ -24,26 +32,26 @@ It of course also makes it easier to store the state.
 
 ## Usage
 
+To apply such an optimiser to a whole model, `setup` builds a tree containing any initial
+state for every trainable array. Then at each step, `update` uses this and the gradient
+to adjust the model:
+
 ```julia
 
 using Flux, Metalhead, Optimisers
 
 o = Optimisers.ADAM() # define an ADAM optimiser with default settings
-st = Optimisers.state(o, m)  # initialize the optimiser before using it
+st = Optimisers.setup(o, m)  # initialize the optimiser before using it
 
-model = ResNet() # define a model to train on
+model = ResNet18() # define a model to train on
 ip = rand(Float32, 224, 224, 3, 1) # dummy data
 
 m̄, _ = gradient(model, ip) do m, x # calculate the gradients
-  sum(m(x))
+  sum(m(x)) # dummy loss function
 end
 
+st, mnew = Optimisers.update(st, m, m̄)
 
-st, mnew = Optimisers.update(o, st, m, m̄)
-
-# or
-
-st, mnew = o(m, m̄, st)
 ```
 
 Notice that a completely new instance of the model is returned. Internally, this
@@ -53,6 +61,6 @@ work with different forms of gradients, but most likely use case are the gradien
 returned by [Zygote.jl](https://fluxml.ai/Zygote.jl).
 
 There is also `Optimisers.update!` which similarly returns a new model and new state,
-but is free to mutate arrays within the old one, for efficiency.
+but is free to mutate arrays within the old one for efficiency.
 The method of `apply!` you write is likewise free to mutate arrays within its state;
 they are defensively copied when this rule is used with `update`.
diff --git a/src/Optimisers.jl b/src/Optimisers.jl
index e70f819..404aabe 100644
--- a/src/Optimisers.jl
+++ b/src/Optimisers.jl
@@ -10,4 +10,97 @@ export Descent, ADAM, Momentum, Nesterov, RMSProp,
        ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, ADAMW, RADAM, OADAM, AdaBelief,
        WeightDecay, ClipGrad, ClipNorm, OptimiserChain
 
+"""
+    Optimisers.apply!(rule::RuleType, state, parameters, gradient) -> (state, gradient)
+
+This defines the action of any optimisation rule. It should return the modified gradient
+which will be subtracted from the parameters, and the updated state (if any) for use at
+the next iteration, as a tuple `(state, gradient)`.
+
+For efficiency it is free to mutate the old state, but only what is returned will be used.
+Ideally this should check `iswriteable(x)`, which the built-in rules do via [`@..`](@ref).
+
+The initial state is `init(rule::RuleType, parameters)`.
+
+# Example
+```jldoctest
+julia> Optimisers.init(Descent(0.1), [1,2,3]) === nothing
+true
+
+julia> Optimisers.apply!(Descent(0.1), nothing, [1,2,3], [4,5,6])
+(nothing, Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}}(*, ([4, 5, 6], 0.1)))
+```
+"""
+apply!
+
+"""
+    Optimisers.init(rule::RuleType, parameters) -> state
+
+Sets up the initial state for a given optimisation rule, and an array of parameters.
+This and [`apply!`](@ref) are the two functions which any new optimisation rule must define.
+
+# Examples
+```jldoctest
+julia> Optimisers.init(Descent(), [1,2,3])  # is `nothing`
+
+julia> Optimisers.init(Momentum(), [1.0, 2.0])
+2-element Vector{Float64}:
+ 0.0
+ 0.0
+```
+"""
+init
+
+"""
+    Optimisers.setup(rule, model) -> tree
+
+Initialises the given optimiser for every trainable parameter within the model.
+Returns a tree of the relevant states, which must be passed to [`update`](@ref)
+or [`update!`](@ref).
+
+# Example
+```jldoctest
+julia> Optimisers.setup(Descent(0.1f0), (x = rand(3), y = (true, false), z = tanh))
+(x = Leaf(Descent{Float32}(0.1), nothing), y = (nothing, nothing), z = nothing)
+```
+"""
+setup
+
+"""
+    Optimisers.update(tree, model, gradient) -> (tree, model)
+
+Uses the optimiser and the gradient to change the trainable parameters in the model.
+Returns the improved model, and the optimiser states needed for the next update.
+The initial tree of states comes from [`setup`](@ref).
+
+See also [`update!`](@ref), which will be faster for models of ordinary `Array`s or `CuArray`s.
+
+# Example
+```jldoctest
+julia> m = (x = Float32[1,2,3], y = tanh);
+
+julia> t = Optimisers.setup(Descent(0.1f0), m)
+(x = Leaf(Descent{Float32}(0.1), nothing), y = nothing)
+
+julia> g = (x = [1,1,1], y = nothing);  # fake gradient
+
+julia> Optimisers.update(t, m, g)
+((x = Leaf(Descent{Float32}(0.1), nothing), y = nothing), (x = Float32[0.9, 1.9, 2.9], y = tanh))
+```
+"""
+update
+
+"""
+    Optimisers.update!(tree, model, gradient) -> (tree, model)
+
+Uses the optimiser and the gradient to change the trainable parameters in the model.
+Returns the improved model, and the optimiser states needed for the next update.
+The initial tree of states comes from [`setup`](@ref).
+
+This is used in exactly the same manner as [`update`](@ref), but because it may mutate
+arrays within the old model (and the old state), it will be faster for models of ordinary
+`Array`s or `CuArray`s. However, you should not rely on the old model being fully updated.
+"""
+update!
+
 end # module
diff --git a/src/interface.jl b/src/interface.jl
index d38c7b7..8a39e04 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -1,39 +1,48 @@
 
-function state(o, x)
+struct Leaf{R,S}
+  rule::R
+  state::S
+end
+
+function setup(rule, x; seen = Base.IdSet())
   if isnumeric(x)
-    return init(o, x)
+    x in seen && throw(ArgumentError("Optimisers.jl does not at present handle tied weights, sorry."))
+    isbits(x) || push!(seen, x)
+    return Leaf(rule, init(rule, x))
   elseif isleaf(x)
     return nothing
   else
     x′, _ = functor(x)
-    return map(xᵢ -> state(o, xᵢ), x′)
+    return map(xᵢ -> setup(rule, xᵢ; seen), x′)
   end
 end
 
-patch!(x, x̄) = iswriteable(x) ? (x .= x .- x̄) : (x .- x̄)
+subtract!(x, x̄) = iswriteable(x) ? (x .= x .- x̄) : (x .- x̄)
 
-function _update!(o, st, x, x̄s...)
-  st′, x̄′ = apply!(o, st, x, x̄s...)
-  return st′, patch!(x, x̄′)
+function update!(ℓ::Leaf, x, x̄s...)
+  if all(isnothing, x̄s)
+    return ℓ, x
+  else
+    s′, x̄′ = apply!(ℓ.rule, ℓ.state, x, x̄s...)
+    return Leaf(ℓ.rule, s′), subtract!(x, x̄′)
+  end
 end
 
-function update!(o, state, x, x̄s...)
+function update!(tree, x, x̄s...)
   if all(isnothing, x̄s)
-    return state, x
-  elseif isnumeric(x)
-    return _update!(o, state, x, x̄s...)
+    return tree, x
   else
     x̄s′ = map(x̄ -> functor(typeof(x), x̄)[1], x̄s)
     x′, re = functor(typeof(x), x)
-    xstate = map((stᵢ, xᵢ, x̄sᵢ...) -> update!(o, stᵢ, xᵢ, x̄sᵢ...), state, x′, x̄s′...)
-    return map(first, xstate), re(map(last, xstate))
+    xtree = map((stᵢ, xᵢ, x̄sᵢ...) -> update!(stᵢ, xᵢ, x̄sᵢ...), tree, x′, x̄s′...)
+    return map(first, xtree), re(map(last, xtree))
   end
 end
 
-function update(o, state, x, x̄s...)
-  state′ = fmap(copy, state; exclude = iswriteable)
+function update(tree, x, x̄s...)
+  t′ = fmap(copy, tree; exclude = iswriteable)
   x′ = fmap(copy, x; exclude = iswriteable)
-  update!(o, state′, x′, x̄s...)
+  update!(t′, x′, x̄s...)
 end
 
 # default all rules to first order calls
@@ -75,3 +84,11 @@ function lazy end
 Broadcast.broadcasted(::typeof(lazy), x) = Lazy(x)
 struct Lazy{T}; bc::T; end
 Broadcast.materialize(x::Lazy) = Broadcast.instantiate(x.bc)
+
+function Base.show(io::IO, ℓ::Leaf)  # show method is mostly to hide its long type!
+  ioc = IOContext(io, :compact => true)
+  print(ioc, "Leaf(", ℓ.rule, ", ")
+  show(ioc, ℓ.state)
+  print(io, ")")
+end
+
diff --git a/src/rules.jl b/src/rules.jl
index 7fe0506..6123e94 100644
--- a/src/rules.jl
+++ b/src/rules.jl
@@ -21,8 +21,6 @@ function apply!(o::Descent, state, x, dx)
   return state, @.. dx * η
 end
 
-(o::Descent)(state, m, dm) = update(o, state, m, dm)
-
 """
     Momentum(η = 1f-2, ρ = 9f-1)
 
@@ -49,8 +47,6 @@ function apply!(o::Momentum, state, x, dx)
   return v′, @.. -v′
 end
 
-(o::Momentum)(state, m, dm) = update(o, state, m, dm)
-
 """
     Nesterov(η = 1f-3, ρ = 9f-1)
 
@@ -70,8 +66,6 @@ Nesterov(η = 1f-3, ρ = 9f-1) = Nesterov{typeof(η)}(η, ρ)
 
 init(o::Nesterov, x::AbstractArray) = zero(x)
 
-(o::Nesterov)(state, m, dm) = update(o, state, m, dm)
-
 function apply!(o::Nesterov, state, x, dx)
   η, ρ, v = o.eta, o.rho, state
   d = @.. ρ^2 * v - (1+ρ) * η * dx
@@ -113,8 +107,6 @@ function apply!(o::RMSProp, state, x, dx)
   return acc′, dx′
 end
 
-(o::RMSProp)(state, m, dm) = update(o, state, m, dm)
-
 """
     ADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))
 
@@ -137,8 +129,6 @@ ADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = ADAM{typeof(η)}(
 
 init(o::ADAM, x::AbstractArray) = (zero(x), zero(x), o.beta)
 
-(o::ADAM)(state, m, dm) = update(o, state, m, dm)
-
 function apply!(o::ADAM{T}, state, x, dx) where T
   η, β, ϵ = o.eta, o.beta, o.epsilon
   mt, vt, βt = state
@@ -172,8 +162,6 @@ RADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = RADAM{typeof(η)}
 
 init(o::RADAM, x::AbstractArray) = (zero(x), zero(x), o.beta, 1)
 
-(o::RADAM)(state, m, dm) = update(o, state, m, dm)
-
 function apply!(o::RADAM, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
   ρ∞ = 2/(1-β[2])-1
@@ -215,8 +203,6 @@ AdaMax(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = AdaMax{typeof(η
 
 init(o::AdaMax, x::AbstractArray) = (zero(x), zero(x), o.beta)
 
-(o::AdaMax)(state, m, dm) = update(o, state, m, dm)
-
 function apply!(o::AdaMax, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
@@ -252,8 +238,6 @@ OADAM(η = 1f-3, β = (5f-1, 9f-1), ϵ = eps(typeof(η))) = OADAM{typeof(η)}(η
 
 init(o::OADAM, x::AbstractArray) = (zero(x), zero(x), o.beta, zero(x))
 
-(o::OADAM)(state, m, dm) = update(o, state, m, dm)
-
 function apply!(o::OADAM, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
@@ -289,8 +273,6 @@ ADAGrad(η = 1f-1, ϵ = eps(typeof(η))) = ADAGrad{typeof(η)}(η, ϵ)
 
 init(o::ADAGrad, x::AbstractArray) = fill!(similar(x), o.epsilon)
 
-(o::ADAGrad)(state, m, dm) = update(o, state, m, dm)
-
 function apply!(o::ADAGrad, state, x, dx)
   η, ϵ = o.eta, o.epsilon
   acc = state
@@ -321,8 +303,6 @@ ADADelta(ρ = 9f-1, ϵ = eps(typeof(ρ))) = ADADelta{typeof(ρ)}(ρ, ϵ)
 
 init(o::ADADelta, x::AbstractArray) = (zero(x), zero(x))
 
-(o::ADADelta)(state, m, dm) = update(o, state, m, dm)
-
 function apply!(o::ADADelta, state, x, dx)
   ρ, ϵ = o.rho, o.epsilon
   acc, Δacc = state
@@ -360,8 +340,6 @@ AMSGrad(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = AMSGrad{typeof(
 init(o::AMSGrad, x::AbstractArray) =
   (fill!(similar(x), o.epsilon), fill!(similar(x), o.epsilon), fill!(similar(x), o.epsilon))
 
-(o::AMSGrad)(state, m, dm) = update(o, state, m, dm)
-
 function apply!(o::AMSGrad, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
@@ -398,8 +376,6 @@ NADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = NADAM{typeof(η)}
 
 init(o::NADAM, x::AbstractArray) = (zero(x), zero(x), o.beta)
 
-(o::NADAM)(state, m, dm) = update(o, state, m, dm)
-
 function apply!(o::NADAM, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
@@ -454,8 +430,6 @@ AdaBelief(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = AdaBelief{typ
 
 init(o::AdaBelief, x::AbstractArray) = (zero(x), zero(x))
 
-(o::AdaBelief)(state, m, dm) = update(o, state, m, dm)
-
 function apply!(o::AdaBelief, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
   mt, st = state
@@ -482,8 +456,6 @@ WeightDecay() = WeightDecay(5f-4)
 
 init(o::WeightDecay, x::AbstractArray) = nothing
 
-(o::WeightDecay)(state, m, dm) = update(o, state, m, dm)
-
 function apply!(o::WeightDecay, state, x, dx)
   dx′ = @.. dx + o.wd * x
 
@@ -504,8 +476,6 @@ ClipGrad() = ClipGrad(10f0)
 
 init(o::ClipGrad, x::AbstractArray) = nothing
 
-(o::ClipGrad)(state::Nothing, m, dm) = update(o, state, m, dm)
-
 function apply!(o::ClipGrad, state, x, dx)
   δ = convert(float(eltype(dx)), o.delta)
   dx′ = @.. clamp(dx, -δ, δ)
@@ -533,8 +503,6 @@ ClipNorm(ω = 10f0, p = 2; throw::Bool = true) = ClipNorm{typeof(ω)}(ω, p, thr
 
 init(o::ClipNorm, x::AbstractArray) = nothing
 
-(o::ClipNorm)(state::Nothing, m, dm) = update(o, state, m, dm)
-
 function apply!(o::ClipNorm, state, x, dx)
   nrm = norm(dx, o.p)
   if o.throw && !isfinite(nrm)
@@ -558,8 +526,6 @@ OptimiserChain(opts...) = OptimiserChain(opts)
 
 init(o::OptimiserChain, x::AbstractArray) = [init(opt, x) for opt in o.opts]
 
-(o::OptimiserChain)(state, m, dms...) = update(o, state, m, dms...)
-
 function apply!(o::OptimiserChain, states, x, dx, dxs...)
   new_states = similar(states)
   for (i, (opt, state)) in enumerate(zip(o.opts, states))
@@ -568,13 +534,3 @@ function apply!(o::OptimiserChain, states, x, dx, dxs...)
 
   return new_states, dx
 end
-
-for Opt in (:Descent, :ADAM, :Momentum, :Nesterov, :RMSProp,
-            :ADAGrad, :AdaMax, :ADADelta, :AMSGrad, :NADAM,
-            :RADAM, :OADAM, :AdaBelief)
-  @eval function $Opt(m::$Opt; kwargs...)
-    fs = fieldnames($Opt)
-    args = NamedTuple{fs}(f in keys(kwargs) ? getindex(kwargs, f) : getfield(m, f) for f in fs)
-    $Opt(args...)
-  end
-end


Test patch:
diff --git a/test/runtests.jl b/test/runtests.jl
index 573a85d..b879d5f 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -11,31 +11,33 @@ using Optimisers: @..
     mid = objectid(m[1])
     g = ([25, 33],)
     o = Descent(0.1)
-    s = Optimisers.state(o, m)
+    s = Optimisers.setup(o, m)
     
-    s2, m2 = Optimisers.update(o, s, m, g)
+    s2, m2 = Optimisers.update(s, m, g)
     @test m[1] == 1:2  # not mutated
     @test Optimisers.iswriteable(m[1])
     @test m2[1] ≈ [1,2] .- 0.1 .* [25, 33]
 
-    s3, m3 = Optimisers.update!(o, s, m, g)
+    s3, m3 = Optimisers.update!(s, m, g)
     @test objectid(m3[1]) == mid
     @test m3[1] ≈ [1,2] .- 0.1 .* [25, 33]
   end
 
-  @testset for o in (Descent(), ADAM(), Momentum(), Nesterov(), RMSProp(),
+  @testset "$(first(string(o), 42))" for o in (
+                     Descent(), ADAM(), Momentum(), Nesterov(), RMSProp(),
                      ADAGrad(), AdaMax(), ADADelta(), AMSGrad(), NADAM(),
-                     ADAMW(), RADAM(), OADAM(), AdaBelief())
+                     ADAMW(), RADAM(), OADAM(), AdaBelief()
+                     )
     w′ = (α = rand(3, 3), β = rand(3, 3))
 
     # Original example
     w = (α = 5rand(3, 3), β = rand(3, 3))
-    st = Optimisers.state(o, w)
+    st = Optimisers.setup(o, w)
     loss(x, y) = mean((x.α .* x.β .- y.α .* y.β) .^ 2)
     @test loss(w, w′) > 1
     for i = 1:10^4
       gs = gradient(x -> loss(x, w′), w)
-      st, w = Optimisers.update(o, st, w, gs...)
+      st, w = Optimisers.update(st, w, gs...)
     end
     lw = loss(w, w′)
     if o isa ADADelta
@@ -46,11 +48,11 @@ using Optimisers: @..
 
     # Slightly harder variant
     m = (α = randn(3), β = transpose(5rand(3,3)), γ = (rand(2), tanh))  # issue 28
-    st = Optimisers.state(o, m)
+    st = Optimisers.setup(o, m)
     @test loss(m, w′) > 1
     for i = 1:10^4
       gs = gradient(x -> loss(x, w′), m)
-      st, m = o(st, m, gs...)
+      st, m = Optimisers.update!(st, m, gs...)
     end
     lm = loss(m, w′)
     if lm < 0.1
@@ -72,38 +74,29 @@ using Optimisers: @..
     for t = 1:10^5
       x = rand(10)
       gs = gradient(w -> loss(x, w, w′), w)
-      st, w = Optimisers.update(opt, st, w, gs...)
+      st, w = Optimisers.update!(st, w, gs...)
     end
-    @test loss(rand(10, 10), w, w′) < 0.01
+    @test_broken loss(rand(10, 10), w, w′) < 0.01
   end
 
   @testset "gradient clipping" begin
-    @test_skip m = (α = ([0], sin), γ = rand(3))  # https://github.com/FluxML/Optimisers.jl/issues/28
-    m = (α = ([0], [0]), γ = rand(3))
-    c1 = ClipGrad(13)
-    s1 = Optimisers.state(c1, m)
-    _, g1 = Optimisers.update(c1, s1, m, (α = nothing, γ = [1,10,100],))
-    @test m.γ .- g1.γ ≈ [1, 10, 13]
+    m = (α = ([0], sin), γ = rand(3))
+    s1 = Optimisers.setup(ClipGrad(13), m)
+    _, m1 = Optimisers.update(s1, m, (α = nothing, γ = [1,10,100],))
+    @test m.γ .- m1.γ ≈ [1, 10, 13]
 
-    c2 = ClipNorm(10)
-    s2 = Optimisers.state(c2, m)
-    _, g2 = Optimisers.update(c2, s2, m, (α = ([0.1], nothing), γ = [1,10,100],))
-    @test only(m.α[1] .- g2.α[1]) ≈ 0.1
-    @test norm(m.γ .- g2.γ) ≈ 10
-    @test_throws DomainError Optimisers.update(c2, s2, m, (α = [0.1], γ = [1,10,NaN],))
+    s2 = Optimisers.setup(ClipNorm(10), m)
+    _, m2 = Optimisers.update(s2, m, (α = ([0.1], nothing), γ = [1,10,100],))
+    @test only(m.α[1] .- m2.α[1]) ≈ 0.1
+    @test norm(m.γ .- m2.γ) ≈ 10
+    @test_throws DomainError Optimisers.update(s2, m, (α = [0.1], γ = [1,10,NaN],))
 
-    c3 = ClipNorm(5, 1; throw=false)
-    _, g3 = Optimisers.update(c3, s2, m, (α = ([0.1], nothing), γ = [1,10,100],))
-    @test only(m.α[1] .- g3.α[1]) ≈ 0.1
-    @test norm(m.γ .- g3.γ, 1) ≈ 5
-    _, g3n = Optimisers.update(c3, s2, m, (α = nothing, γ = [1,10,Inf],))
-    @test isnan(g3n.γ[3])
-  end
-
-  @testset "Optimiser Updates" begin
-    opt = ADAM()
-    new_opt = ADAM(opt, eta = 9.f0)
-    @test new_opt.eta == 9.f0
+    s3 = Optimisers.setup(ClipNorm(5, 1; throw=false), m)
+    _, m3 = Optimisers.update(s3, m, (α = ([0.1], nothing), γ = [1,10,100],))
+    @test only(m.α[1] .- m3.α[1]) ≈ 0.1
+    @test norm(m.γ .- m3.γ, 1) ≈ 5
+    _, m3n = Optimisers.update!(s3, m, (α = nothing, γ = [1,10,Inf],))
+    @test isnan(m3n.γ[3])
   end
 
   @testset "broadcasting macro" begin
@@ -119,4 +112,13 @@ using Optimisers: @..
     @test (@.. r = y * z) == y .* z
   end
 
+  @testset "tied weights" begin
+    ok = (1.0:3.0, sin, "abc", :abc)
+    m = (α = ok, β = rand(3), γ = ok)
+    m1 = (rand(3), m, rand(3))
+    @test Optimisers.setup(ADAMW(), m1) isa Tuple
+    m2 = (rand(3), m, rand(3), m, rand(3))  # illegal
+    @test_throws ArgumentError Optimisers.setup(ADAMW(), m2)
+  end
+
 end