Move all optimizers to Optimisers.jl

src/Optimisers.jl

@@ -5,7 +5,9 @@ using Functors: functor, fmap, isleaf
 include("interface.jl")
 include("rules.jl")
 
-export Descent, Momentum, Nesterov, RMSProp,
-       ADAM
+export init, update!
+export Descent, ADAM, Momentum, Nesterov, RMSProp,
+       ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, ADAMW, RADAM, OADAM, AdaBelief,
+       WeightDecay, Sequence
 
 end # module

src/interface.jl

@@ -1,30 +1,31 @@
-function patch(x, x̄)
-  return x .- x̄
+function patch!(x, x̄)
+  @. x .-= x̄
+  return x
 end
 
-function state(o, x)
+function init(o, x)
   if isleaf(x)
     return init(o, x)
   else
     x, _ = functor(x)
-    map(x -> state(o, x), x)
+    return map(x -> init(o, x), x)
   end
 end
 
-function _update(o, x, x̄, st)
-  x̄, st = apply(o, x, x̄, st)
-  return patch(x, x̄), st
+function _update!(o, x, x̄, st)
+  x̄, st = apply!(o, x, x̄, st)
+  return patch!(x, x̄), st
 end
 
-function update(o, x::T, x̄, state) where T
+function update!(o, x::T, x̄, state) where T
   if x̄ === nothing
     return x, state
   elseif isleaf(x)
-    _update(o, x, x̄, state)
+    _update!(o, x, x̄, state)
   else
     x̄, _  = functor(typeof(x), x̄)
-    x, re = functor(typeof(x), x)
-    xstate = map((x, x̄, state) -> update(o, x, x̄, state), x, x̄, state)
-    re(map(first, xstate)), map(x -> x[2], xstate)
+    x, restructure = functor(typeof(x), x)
+    xstate = map((x, x̄, state) -> update!(o, x, x̄, state), x, x̄, state)
+    restructure(map(first, xstate)), map(x -> x[2], xstate)
   end
 end