Repository: FluxML/Optimisers.jl
Pull request #31 (2022-01-27 21:48:49)

Add lazy broadcasting
This lets the each optimiser step build a `Broadcasted` object, which is ultimately evaluated once, in the usual way. They should never escape, but they do save memory:
```
julia> m = (α = rand(100,100), β = rand(100,100));

julia> o = OptimiserChain(WeightDecay(), ADAM());

julia> st = Optimisers.state(o, m);

julia> gr = (α = rand(100,100), β = rand(100,100)); # fake gradient

julia> @btime Optimisers.update($o, $st, $m, $gr);
  min 115.250 μs, mean 180.870 μs (55 allocations, 470.42 KiB)  # before
  min 108.875 μs, mean 126.485 μs (46 allocations, 158.27 KiB)  # used for optimisers
  min 43.750 μs, mean 47.868 μs (46 allocations, 3.86 KiB)      # with in-place model update
```
The same macro also inserts a check that it is possible for optimisers to write into their own state arrays, before doing so. Right now they always try, and they make arrays like `zero(x)`, so will fail for e.g. StaticArrays.

We could similarly use that to mutate the arrays of the model, when possible. But perhaps that wants to be optional even if it's the default, maybe it's `update(opt, state, model, grad; copy=false)`?

It's possible that this option ought to propagate down to every rule, to also control whether they mutate their state arrays. That would mean passing it to every `apply` method, which seems messy.

darsnack: If we are going to "maybe mutate" then I suggest appending `!` to the function names to signal that.
mcabbott: Having `update!` and `update` would be fine as an interface, but I'm keen to avoid having `apply!` and `apply` for every optimiser. There should only be one place you can make sign errors. 
darsnack: I agree, there should only be `apply!`.
darsnack: Would it be slightly less spooky if we use something global to the package like `use_cuda` in Flux? Having a dangling keyword argument in the interface is not so great for people wanting to write their own rules.
mcabbott: So the last commit makes the interface `update` fully out of place, and `update!` fully overwriting (when possible). 

It imposes on anyone else writing a rule that they must accept this keyword (and ideally obey it), but nothing too weird will happen -- you will get an ordinary MethodError if your new `apply` cannot be called, same as if you made other mistakes.
ToucheSir: My thought in https://github.com/FluxML/Optimisers.jl/pull/31#discussion_r794057657 was to have rules always try mutation, then provide a higher-level interface that defensively copies any mutable arrays in state/gradients before calling `update!` and `apply!`. Thoughts?
darsnack: I would prefer that approach or checking `Optimisers.use_mutation[]` in the macro to the magic local keyword. Just a little too magical for me.
mcabbott: For defensive copies, I guess it's a reasonable assumption that any optimiser planned to mutate any array it put in its state. What you do need is a way to access them, i.e. you need to assume that some `fmap(maybecopy, state)` is going to work. Which you can document. The typical `apply!` will then contain things like `v′ = @.. v = ρ * v` with the (less magical) macro mutating v if mutable -- with the same source of truth as the defensive `maybecopy` process.
ToucheSir: I skimmed the stdlib, StaticArrays and CUDA.jl. All three seem to have functioning `copy` functions, so here's a sketch of `update`:
```julia
maybecopy(x) = iswriteable(x) ? copy(x) : x
update(o, state, x, x̄s...) = update!(o, fmap(maybecopy, state, exclude=isnumeric), x, fmap(maybecopy, x̄s, exclude=isnumeric)...)
```
  
mcabbott: Oh that's another way. Maybe tidier than doing it at the `_update!` stage which acts on one trainable array, and you only pay for one `fmap` cache...

We could go further and simply demand that the state be a tuple or a NamedTuple, always, and call `map`. That might be less mysterious. E.g. if your state has some struct and you forgot to `@functor` it, then these other ways may silently give you unintentional mutation... but `map` will loudly fail the first time you try.
ToucheSir: For better or worse, `fmap` is still needed for handling nested namedtuples. Un-functored structs should be safe as long as we guard with `isleaf(x) && !isnumeric(x)` (the default definition of `functor` returns an empty tuple of children for this reason). 
mcabbott: Maybe that's fine, `update` is going to allocate so much that an `IdDict` is really no concern. I've adopted this `fmap` strategy. Maybe it's even a good thing if it walks not just the trainable arrays (once we restore that distinction). But `fmap` with `exclude` isn't in the tagged version?

And I updated all the rules. Besides tests it might be worth someone else looking over them in case there are mistakes.

Patch:
diff --git a/Manifest.toml b/Manifest.toml
deleted file mode 100644
index 23989dc..0000000
--- a/Manifest.toml
+++ /dev/null
@@ -1,57 +0,0 @@
-# This file is machine-generated - editing it directly is not advised
-
-[[Base64]]
-uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
-
-[[Functors]]
-deps = ["MacroTools"]
-git-tree-sha1 = "bd96c1974893293ecc8b7e9016954d8cbd840ee4"
-repo-rev = "master"
-repo-url = "https://github.com/FluxML/Functors.jl"
-uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
-version = "0.1.0"
-
-[[Libdl]]
-uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
-
-[[LinearAlgebra]]
-deps = ["Libdl"]
-uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-
-[[MacroTools]]
-deps = ["Markdown", "Random"]
-git-tree-sha1 = "6a8a2a625ab0dea913aba95c11370589e0239ff0"
-uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
-version = "0.5.6"
-
-[[Markdown]]
-deps = ["Base64"]
-uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
-
-[[Random]]
-deps = ["Serialization"]
-uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-
-[[Requires]]
-deps = ["UUIDs"]
-git-tree-sha1 = "cfbac6c1ed70c002ec6361e7fd334f02820d6419"
-uuid = "ae029012-a4dd-5104-9daa-d747884805df"
-version = "1.1.2"
-
-[[SHA]]
-uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
-
-[[Serialization]]
-uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
-
-[[SparseArrays]]
-deps = ["LinearAlgebra", "Random"]
-uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-
-[[Statistics]]
-deps = ["LinearAlgebra", "SparseArrays"]
-uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
-
-[[UUIDs]]
-deps = ["Random", "SHA"]
-uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
diff --git a/Project.toml b/Project.toml
index a12b047..a688310 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,19 +1,17 @@
 name = "Optimisers"
 uuid = "3bd65402-5787-11e9-1adc-39752487f4e2"
 authors = ["Mike J Innes <mike.j.innes@gmail.com>"]
-version = "0.1.0"
+version = "0.2.0"
 
 [deps]
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-Functors = "0.1, 0.2"
-Requires = "0.5, 1"
-julia = "1.6" 
+Functors = "0.2.7"
+julia = "1.6"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/docs/src/api.md b/docs/src/api.md
index fdcf2fc..8ef3a33 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -15,7 +15,7 @@ Optimisers.ADAMW
 Optimisers.AdaBelief
 ```
 
-```
+```@docs
 Optimisers.ClipGrad
 Optimisers.ClipNorm
 Optimisers.WeightDecay
diff --git a/docs/src/index.md b/docs/src/index.md
index b2d7cbd..c0d23ac 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -8,9 +8,9 @@ struct Descent{T}
   η::T
 end
 
-# Define an `apply` rule with which to update the current params
+# Define an `apply!` rule with which to update the current params
 # using the gradients
-function Optimisers.apply(o::Descent, state, m, m̄)
+function Optimisers.apply!(o::Descent, state, m, m̄)
     o.η .* m̄, state
 end
 
@@ -52,3 +52,7 @@ tree formed by the model and update the parameters using the gradients. Optimise
 work with different forms of gradients, but most likely use case are the gradients as
 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.
+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/interface.jl b/src/interface.jl
index 7768b85..d38c7b7 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -1,4 +1,3 @@
-patch(x, x̄) = x .- x̄
 
 function state(o, x)
   if isnumeric(x)
@@ -11,27 +10,68 @@ function state(o, x)
   end
 end
 
-function _update(o, st, x, x̄s...)
-  st′, x̄′ = apply(o, st, x, x̄s...)
-  return st′, patch(x, x̄′)
+patch!(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̄′)
 end
 
-function update(o, state, x::T, x̄s...) where T
+function update!(o, state, x, x̄s...)
   if all(isnothing, x̄s)
     return state, x
   elseif isnumeric(x)
-    return _update(o, state, x, x̄s...)
+    return _update!(o, state, x, x̄s...)
   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′...)
+    xstate = map((stᵢ, xᵢ, x̄sᵢ...) -> update!(o, stᵢ, xᵢ, x̄sᵢ...), state, x′, x̄s′...)
     return map(first, xstate), re(map(last, xstate))
   end
 end
 
+function update(o, state, x, x̄s...)
+  state′ = fmap(copy, state; exclude = iswriteable)
+  x′ = fmap(copy, x; exclude = iswriteable)
+  update!(o, state′, x′, x̄s...)
+end
+
 # default all rules to first order calls
-apply(o, state, x, dx, dxs...) = apply(o, state, x, dx)
+apply!(o, state, x, dx, dxs...) = apply!(o, state, x, dx)
 
 isnumeric(x::AbstractArray{<:Number}) = isleaf(x)  # isleaf to allow for e.g. transposed shared weights
 isnumeric(x::AbstractArray{<:Bool}) = false  # convention of ChainRules is that Bool is non-differentiable
 isnumeric(x) = false
+
+iswriteable(::DenseArray{<:AbstractFloat}) = true  # more elaborate versions are possible, wait until needed?
+iswriteable(_) = false
+
+"""
+    @.. x = x + y
+    @.. x + y / z
+
+Magic broadcasting macro, for use in `apply!` rules:
+* Applied to assignment `x = ...` it is like `@.` unless `!iswriteable(x)`,
+  in which case it ignores `x`, and applies `@.` on the right.
+* Applied to other expressions, it broadcasts like `@.` but does not materialise,
+  returning a `Broadcasted` object for later use.
+"""
+macro var".."(ex)
+  if Meta.isexpr(ex, :(=))
+    dst = esc(ex.args[1])
+    src = esc(Broadcast.__dot__(ex.args[2]))
+    :(if $iswriteable($dst)
+        $dst .= $src
+      else
+        $src
+      end)
+  else
+    bc = esc(Broadcast.__dot__(ex))
+    :($lazy.($bc))
+  end
+end
+
+function lazy end
+Broadcast.broadcasted(::typeof(lazy), x) = Lazy(x)
+struct Lazy{T}; bc::T; end
+Broadcast.materialize(x::Lazy) = Broadcast.instantiate(x.bc)
diff --git a/src/rules.jl b/src/rules.jl
index b19c2cd..7fe0506 100644
--- a/src/rules.jl
+++ b/src/rules.jl
@@ -15,10 +15,10 @@ Descent() = Descent(1f-1)
 
 init(o::Descent, x::AbstractArray) = nothing
 
-function apply(o::Descent, state, x, dx)
-  η = convert(eltype(dx), o.eta)
+function apply!(o::Descent, state, x, dx)
+  η = convert(float(eltype(dx)), o.eta)
   
-  return state, dx .* η
+  return state, @.. dx * η
 end
 
 (o::Descent)(state, m, dm) = update(o, state, m, dm)
@@ -42,11 +42,11 @@ Momentum(η = 1f-2, ρ = 9f-1) = Momentum{typeof(η)}(η, ρ)
 
 init(o::Momentum, x::AbstractArray) = zero(x)
 
-function apply(o::Momentum, state, x, dx)
+function apply!(o::Momentum, state, x, dx)
   η, ρ, v = o.eta, o.rho, state
-  @. v = ρ * v - η * dx
+  v′ = @.. v = ρ * v - η * dx
   
-  return v, -v
+  return v′, @.. -v′
 end
 
 (o::Momentum)(state, m, dm) = update(o, state, m, dm)
@@ -72,12 +72,12 @@ init(o::Nesterov, x::AbstractArray) = zero(x)
 
 (o::Nesterov)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::Nesterov, state, x, dx)
+function apply!(o::Nesterov, state, x, dx)
   η, ρ, v = o.eta, o.rho, state
-  d = @. ρ^2 * v - (1+ρ) * η * dx
-  @. v = ρ * v - η * dx
+  d = @.. ρ^2 * v - (1+ρ) * η * dx
+  v′ = @.. v = ρ * v - η * dx
   
-  return v, -d
+  return v′, @.. -d
 end
 
 """
@@ -105,12 +105,12 @@ RMSProp(η = 1f-3, ρ = 9f-1, ϵ = eps(typeof(η))) = RMSProp{typeof(η)}(η, ρ
 
 init(o::RMSProp, x::AbstractArray) = zero(x)
 
-function apply(o::RMSProp, state, x, dx)
+function apply!(o::RMSProp, state, x, dx)
   η, ρ, ϵ, acc = o.eta, o.rho, o.epsilon, state
-  @. acc = ρ * acc + (1 - ρ) * dx^2
-  dx′ = @. dx * (η / (sqrt(acc) + ϵ))
+  acc′ = @.. acc = ρ * acc + (1 - ρ) * dx^2
+  dx′ = @.. dx * (η / (sqrt(acc) + ϵ))
   
-  return acc, dx′
+  return acc′, dx′
 end
 
 (o::RMSProp)(state, m, dm) = update(o, state, m, dm)
@@ -139,15 +139,15 @@ 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
+function apply!(o::ADAM{T}, state, x, dx) where T
   η, β, ϵ = o.eta, o.beta, o.epsilon
   mt, vt, βt = state
 
-  @. mt = β[1] * mt + (one(T) - β[1]) * dx
-  @. vt = β[2] * vt + (one(T) - β[2]) * dx ^ 2
-  dx′ = @. mt / (one(T) - βt[1]) / (sqrt(vt / (one(T) - βt[2])) + ϵ) * η
+  mt′ = @.. mt = β[1] * mt + (one(T) - β[1]) * dx
+  vt′ = @.. vt = β[2] * vt + (one(T) - β[2]) * dx ^ 2
+  dx′ = @.. mt / (one(T) - βt[1]) / (sqrt(vt / (one(T) - βt[2])) + ϵ) * η
 
-  return (mt, vt, βt .* β), dx′
+  return (mt′, vt′, βt .* β), dx′
 end
 
 """
@@ -174,23 +174,23 @@ 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)
+function apply!(o::RADAM, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
   ρ∞ = 2/(1-β[2])-1
 
   mt, vt, βt, t = state
 
-  @. mt = β[1] * mt + (1 - β[1]) * dx
-  @. vt = β[2] * vt + (1 - β[2]) * dx^2
+  mt′ = @.. mt = β[1] * mt + (1 - β[1]) * dx
+  vt′ = @.. vt = β[2] * vt + (1 - β[2]) * dx^2
   ρ = ρ∞ - 2*t * βt[2] / (1 - βt[2])
   if ρ > 4
     r = sqrt((ρ - 4) * (ρ - 2) * ρ∞/((ρ∞ - 4) * (ρ∞ - 2) * ρ))
-    dx′ = @. mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ) * η * r
+    dx′ = @.. mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ) * η * r
   else
-    dx′ = @. mt / (1 - βt[1]) * η
+    dx′ = @.. mt / (1 - βt[1]) * η
   end
 
-  return (mt, vt, βt .* β, t + 1), dx′
+  return (mt′, vt′, βt .* β, t + 1), dx′
 end
 
 """
@@ -217,16 +217,16 @@ 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)
+function apply!(o::AdaMax, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
   mt, ut, βt = state
 
-  @. mt = β[1] * mt + (1 - β[1]) * dx
-  @. ut = max(β[2] * ut, abs(dx))
-  dx′ = @. (η/(1 - βt[1])) * mt/(ut + ϵ)
+  mt′ = @.. mt = β[1] * mt + (1 - β[1]) * dx
+  ut′ = @.. ut = max(β[2] * ut, abs(dx))
+  dx′ = @.. (η/(1 - βt[1])) * mt/(ut + ϵ)
 
-  return (mt, ut, βt .* β), dx′
+  return (mt′, ut′, βt .* β), dx′
 end
 
 """
@@ -254,18 +254,18 @@ 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)
+function apply!(o::OADAM, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
   mt, vt, βt, dx_ = state
 
-  @. mt = β[1] * mt + (1 - β[1]) * dx
-  @. vt = β[2] * vt + (1 - β[2]) * dx^2
-  @. dx = -dx_
-  @. dx_ = η * mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ)
-  dx′ = @. dx + 2*dx_
+  mt′ = @.. mt = β[1] * mt + (1 - β[1]) * dx
+  vt′ = @.. vt = β[2] * vt + (1 - β[2]) * dx^2
+  dx = @.. -dx_
+  dx_′ = @.. dx_ = η * mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ)
+  dx′ = @.. dx + 2*dx_
 
-  return (mt, vt, βt .* β, dx_), dx′
+  return (mt′, vt′, βt .* β, dx_′), dx′
 end
 
 """
@@ -291,14 +291,14 @@ 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)
+function apply!(o::ADAGrad, state, x, dx)
   η, ϵ = o.eta, o.epsilon
   acc = state
 
-  @. acc += dx^2
-  dx′ = @. dx * η / (sqrt(acc) + ϵ)
+  acc′ = @.. acc = acc + dx^2
+  dx′ = @.. dx * η / (sqrt(acc) + ϵ)
 
-  return acc, dx′
+  return acc′, dx′
 end
 
 """
@@ -323,17 +323,17 @@ 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)
+function apply!(o::ADADelta, state, x, dx)
   ρ, ϵ = o.rho, o.epsilon
   acc, Δacc = state
 
-  @. acc = ρ * acc + (1 - ρ) * dx^2
+  acc′ = @.. acc = ρ * acc + (1 - ρ) * dx^2
   # DON'T remove epsilon from numerator
   # or even out of the square roots
-  dx′ = @. dx * sqrt(Δacc + ϵ) / sqrt(acc + ϵ)
-  @. Δacc = ρ * Δacc + (1 - ρ) * dx^2
+  dx′ = @.. dx * sqrt(Δacc + ϵ) / sqrt(acc + ϵ)
+  Δacc′ = @.. Δacc = ρ * Δacc + (1 - ρ) * dx^2
   
-  return (acc, Δacc), dx′
+  return (acc′, Δacc′), dx′
 end
 
 """
@@ -362,17 +362,17 @@ init(o::AMSGrad, x::AbstractArray) =
 
 (o::AMSGrad)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::AMSGrad, state, x, dx)
+function apply!(o::AMSGrad, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
   mt, vt, v̂t = state
 
-  @. mt = β[1] * mt + (1 - β[1]) * dx
-  @. vt = β[2] * vt + (1 - β[2]) * dx ^ 2
-  @. v̂t = max(v̂t, vt)
-  dx′ = @. η * mt / (sqrt(v̂t) + ϵ)
+  mt′ = @.. mt = β[1] * mt + (1 - β[1]) * dx
+  vt′ = @.. vt = β[2] * vt + (1 - β[2]) * dx ^ 2
+  v̂t′ = @.. v̂t = max(v̂t, vt)
+  dx′ = @.. η * mt / (sqrt(v̂t) + ϵ)
 
-  return (mt, vt, v̂t), dx′
+  return (mt′, vt′, v̂t′), dx′
 end
 
 """
@@ -400,17 +400,17 @@ 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)
+function apply!(o::NADAM, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
   mt, vt, βt = state
 
-  @. mt = β[1] * mt + (1 - β[1]) * dx
-  @. vt = β[2] * vt + (1 - β[2]) * dx^2
-  dx′ = @. (β[1] * mt / (1 - β[1] * βt[1]) + (1 - β[1]) * dx / (1 - βt[1])) / 
+  mt′ = @.. mt = β[1] * mt + (1 - β[1]) * dx
+  vt′ = @.. vt = β[2] * vt + (1 - β[2]) * dx^2
+  dx′ = @.. (β[1] * mt / (1 - β[1] * βt[1]) + (1 - β[1]) * dx / (1 - βt[1])) / 
           (sqrt(vt * β[2] / (1 - βt[2])) + ϵ) * η
 
-  return (mt, vt, βt .* β), dx′
+  return (mt′, vt′, βt .* β), dx′
 end
 
 """
@@ -456,15 +456,15 @@ 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)
+function apply!(o::AdaBelief, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
   mt, st = state
 
-  @. mt = β[1] * mt + (1 - β[1]) * dx
-  @. st = β[2] * st + (1 - β[2]) * (dx - mt)^2
-  dx′ = @. η * mt / (sqrt(st) + ϵ)
+  mt′ = @.. mt = β[1] * mt + (1 - β[1]) * dx
+  st′ = @.. st = β[2] * st + (1 - β[2]) * (dx - mt)^2
+  dx′ = @.. η * mt / (sqrt(st) + ϵ)
   
-  return (mt, st), dx′
+  return (mt′, st′), dx′
 end
 
 """
@@ -484,8 +484,8 @@ 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
+function apply!(o::WeightDecay, state, x, dx)
+  dx′ = @.. dx + o.wd * x
 
   return state, dx′
 end
@@ -506,9 +506,9 @@ 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(eltype(dx), o.delta)
-  dx′ = @. clamp(dx, -δ, δ)
+function apply!(o::ClipGrad, state, x, dx)
+  δ = convert(float(eltype(dx)), o.delta)
+  dx′ = @.. clamp(dx, -δ, δ)
 
   return state, dx′
 end
@@ -535,14 +535,14 @@ init(o::ClipNorm, x::AbstractArray) = nothing
 
 (o::ClipNorm)(state::Nothing, m, dm) = update(o, state, m, dm)
 
-function apply(o::ClipNorm, state, x, dx)
+function apply!(o::ClipNorm, state, x, dx)
   nrm = norm(dx, o.p)
   if o.throw && !isfinite(nrm)
     throw(DomainError("gradient has $(o.p)-norm $nrm, for array $(summary(x))"))
   end
   λ = min(o.omega / nrm, 1)
 
-  return state, @. dx * λ
+  return state, @.. dx * λ
 end
 
 """
@@ -560,10 +560,10 @@ 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...)
+function apply!(o::OptimiserChain, states, x, dx, dxs...)
   new_states = similar(states)
   for (i, (opt, state)) in enumerate(zip(o.opts, states))
-    new_states[i], dx = apply(opt, state, x, dx, dxs...)
+    new_states[i], dx = apply!(opt, state, x, dx, dxs...)
   end
 
   return new_states, dx


Test patch:
diff --git a/test/runtests.jl b/test/runtests.jl
index 0768dd0..573a85d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,10 +1,28 @@
 using Optimisers, Test
 using Zygote
 using Statistics, Random, LinearAlgebra
-Random.seed!(84)
+Random.seed!(1)
+using Optimisers: @..
 
 @testset verbose=true "Optimisers.jl" begin
 
+  @testset "very basics" begin
+    m = ([1.0, 2.0],)
+    mid = objectid(m[1])
+    g = ([25, 33],)
+    o = Descent(0.1)
+    s = Optimisers.state(o, m)
+    
+    s2, m2 = Optimisers.update(o, 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)
+    @test objectid(m3[1]) == mid
+    @test m3[1] ≈ [1,2] .- 0.1 .* [25, 33]
+  end
+
   @testset for o in (Descent(), ADAM(), Momentum(), Nesterov(), RMSProp(),
                      ADAGrad(), AdaMax(), ADADelta(), AMSGrad(), NADAM(),
                      ADAMW(), RADAM(), OADAM(), AdaBelief())
@@ -20,7 +38,11 @@ Random.seed!(84)
       st, w = Optimisers.update(o, st, w, gs...)
     end
     lw = loss(w, w′)
-    @test lw < 0.001
+    if o isa ADADelta
+      @test_broken lw < 0.001
+    else
+      @test lw < 0.001
+    end
 
     # Slightly harder variant
     m = (α = randn(3), β = transpose(5rand(3,3)), γ = (rand(2), tanh))  # issue 28
@@ -84,4 +106,17 @@ Random.seed!(84)
     @test new_opt.eta == 9.f0
   end
 
+  @testset "broadcasting macro" begin
+    x = [1.0, 2.0]; y = [3,4]; z = [5,6]
+    @test (@.. x + y * z) isa Broadcast.Broadcasted
+    bc = @.. x + y * z
+    @test (y .+ 2 .* bc) == [35,56]
+
+    @test (@.. x = y * z) isa Array
+    @test x == y .* z  # mutated
+    r = 1.0:2.0  # immutable
+    @test (@.. r = y * z) isa Array
+    @test (@.. r = y * z) == y .* z
+  end
+
 end