Repository: FluxML/Optimisers.jl
Pull request #48 (2022-01-30 19:03:05)

Separate `@lazy` from `@..` macro, fix some bugs
This changes the behaviour of `@.. x = rhs` on immutable objects, to attach the label to the new result. While this may mean that the type of `x` changes, it also fixes several bugs where later lines assumed that `x` would be the new one. Should have noticed this on the first pass, sorry. (This was added in #31.)

There are more subtle ways for things to go wrong. Flux's implementation of OADAM uses the mutable gradient array as the place to save an intermediate step: 
https://github.com/FluxML/Flux.jl/blob/05a608b3fb540661bf68c339793cd7083e2ad400/src/optimise/optimisers.jl#L311-L326
That won't work if `Δ == dx` is immutable, but attaching the label to the RHS also won't work, if that gets mutated in the next step. And in general mutating the gradient is unsafe, if could also be someone else's gradient. ~~So for now I make it explicitly materialise the new `dx`, and then update the saved term. Broadcasting a `sqrt` is not super-expensive.~~ Changed now to an explicit copy of `term`, before updating.

<img width="774" alt="Screenshot 2022-01-30 at 13 56 53" src="https://user-images.githubusercontent.com/32575566/151713454-688ecb11-d3b1-4217-bd20-e4fae80d6dfc.png">


In fact, the whole lazy-Broadcasted idea assumes that the function being broadcasted isn't expensive. E.g. `ADADelta` doesn't have correctness bugs, but will compute the sqrt etc. when doing `Δacc =`, and again when `dx′` is ultimately consumed in `subtract!`. 

ToucheSir: On laziness: perhaps `@..` should only handle the `@.. x = ...` case and always materialize. Then laziness could be an explicit opt-in with `lazy.(@. ...)`&mdash;not any more verbose and reduces the trickiness involved.

On ADADelta: is there any reason `dx` can't be materialized into directly instead of creating a lazy `dx'`? Seems like this would be worth the extra kernel dispatch + memory usage given it's re-used during `subtract!`.
mcabbott: I do think we want laziness 99% of the time, so prefer making this compact. It's what lets ClipGrad etc. compose almost for free. 

OADAM actually has a typo which makes it materialise the gradient. It should have `dx′ = @.. 2 * term - prev`. Fusing this 2* and - with `subtract!` seems like what we want, they are free, but a copy is not.

Re ADADelta ... it might be worth some attempt to benchmark this one in isolation?
ToucheSir: Would a different macro or even a short alias for `lazy` help there? If the contention is that having `@.. x = ...` and `@.. ...` return different types and observe different behaviour wrt laziness and materialization can cause subtle issues, perhaps that's a sign that they should not have the same name.
mcabbott: ADADelta times:
```julia
julia> using Optimisers, BenchmarkTools
julia> m, g = ntuple(_ -> rand(Float32, 100,100), 10), ntuple(_ -> rand(Float32, 100,100), 10);
julia> s = Optimisers.setup(ADADelta(), m);

julia> @btime Optimisers.update!($s, $m, $g);
  min 414.583 μs, mean 419.602 μs (0 allocations)  # lazy
  min 334.500 μs, mean 381.109 μs (20 allocations, 391.09 KiB)  # eager

  min 377.166 μs, mean 379.411 μs (0 allocations)  # lazy + @fastmath
  min 180.708 μs, mean 214.607 μs (20 allocations, 391.09 KiB)  # eager + @fastmath

julia> s = Optimisers.setup(ADAM(), m);
julia> @btime Optimisers.update!($s, $m, $g);
  min 265.583 μs, mean 268.664 μs (20 allocations, 320 bytes)

julia> s = Optimisers.setup(Descent(), m);
julia> @btime Optimisers.update!($s, $m, $g);
  min 36.708 μs, mean 37.492 μs (0 allocations)

 julia> m2 = ntuple(_ -> ntuple(_ -> rand(Float32, 1000,1000), 20), 20);
julia> g2 = ntuple(_ -> ntuple(_ -> rand(Float32, 1000,1000), 20), 20);
julia> s2 = Optimisers.setup(ADADelta(), m2);

julia> @btime Optimisers.update!($s2, $m2, $g2);
  min 3.710 s, mean 3.710 s (0 allocations)
  min 3.429 s, mean 3.429 s (800 allocations, 1.49 GiB)
```
On an M1 laptop. I guess checking with CuArrays would be better. 

Re bugs, I think the problem was mostly me adding \prime too many places mechanically without thinking. That I am sure I could have managed even with long descriptive macro names. I will take one more read through this to double-check.
ToucheSir: All numbers are post warmup.

Lazy:
```julia
julia> CUDA.@sync @btime Optimisers.update!($s, $m, $g);
  215.351 μs (2080 allocations: 267.03 KiB)

julia> let s=s, m=m, g=g; CUDA.@time Optimisers.update!(s, m, g); end;
  0.000329 seconds (2.08 k CPU allocations: 267.359 KiB)
```
Eager:
```julia
julia> CUDA.@sync @btime Optimisers.update!($s, $m, $g);
  184.293 μs (1770 allocations: 166.41 KiB)

julia> let s=s, m=m, g=g; CUDA.@time Optimisers.update!(s, m, g); end;
  0.000280 seconds (1.77 k CPU allocations: 166.734 KiB)
```
Big lazy:
```julia
julia> CUDA.@sync @btime Optimisers.update!($s2, $m2, $g2);
  10.340 ms (83200 allocations: 10.43 MiB)

julia> let s=s2, m=m2, g=g2; CUDA.@time Optimisers.update!(s, m, g); end;
  0.055469 seconds (83.22 k CPU allocations: 10.444 MiB)
```
Big eager:
```julia
julia> CUDA.@sync @btime Optimisers.update!($s2, $m2, $g2);
  21.282 ms (70800 allocations: 6.50 MiB)

julia> let s=s2, m=m2, g=g2; CUDA.@time Optimisers.update!(s, m, g); end;
  0.062318 seconds (70.81 k CPU allocations: 6.513 MiB, 5.08% gc time)
```
mcabbott: Sorry, I timed some things, & then deleted them as I think the lazy one gets the wrong answer. It captures Δacc which is mutated on the next line. Maybe Nesterov is wrong too.

Besides proofreading this probably wants bolder notation. I think also some of the `dx′` are pretty subtle, about which terms get the old and the new, I wonder if there too we ought to have something bolder.
ToucheSir: I didn't catch that despite multiple readthroughs as well. How about we stick with eager `@..` for now until to let the design of the lazy path bake? Users who care about memory usage will be using `writeable` arrays anyhow, so it shouldn't create significant allocation overhead.
mcabbott: Here's one test we should have had:
```
@testset verbose=true "mutation check" begin
  @testset "$(name(o))" for o in RULES
    model = Float64.(rand(Int8, 8))
    s_model = SVector{8}(model)
    grads = [Float64.(rand(Int8, 8)) for t in 1:13]
    s_grads = [SVector{8}(x) for x in grads]
    state = Optimisers.setup(o, model)
    s_state = Optimisers.setup(o, s_model)
    for t in 1:13
      state, model = Optimisers.update!(state, model, grads[t])
      s_state, s_model = Optimisers.update!(s_state, s_model, s_grads[t])
    end
    @test model == s_model 
  end
end
```
This would catch the laziness bug, since mutation of an array you have captured doesn't affect the SArray path, and it does in practice. (In the 2nd digit, not the 16th.) It fails for ADAGrad and AMSGrad, at Float32 resolution, because for these two `map(_ -> o.epsilon, x)` makes an array of Float32, not 64.

Can this miss other such bugs? 
mcabbott: Here's a check that results match Flux exactly: https://gist.github.com/mcabbott/886756eb8a8ee3c8cedd5173ded13aa3

Patch:
diff --git a/docs/src/api.md b/docs/src/api.md
index 7e5bcf2..5671140 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -48,4 +48,5 @@ Optimisers.trainable
 Optimisers.apply!
 Optimisers.init
 Optimisers.@..
+Optimisers.@lazy
 ```
diff --git a/src/interface.jl b/src/interface.jl
index 837ae57..6ef29f6 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -78,27 +78,32 @@ end
 
 """
     @.. 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.
+Sometimes in-place broadcasting macro, for use in `apply!` rules.
+If `iswriteable(x)` then it is just `@. x = rhs`, but if not, it becomes `x = @. rhs`.
 """
 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
+  Meta.isexpr(ex, :(=)) || throw("the macro @.. only accepts assignment, like @.. x = y + z")
+  dst = esc(ex.args[1])
+  src = esc(Broadcast.__dot__(ex.args[2]))
+  :($dst = if $iswriteable($dst)
+      $dst .= $src
+    else
+      $src
+    end)
+end
+
+"""
+    x = @lazy y + z
+
+Lazy broadcasting macro, for use in `apply!` rules. It broadcasts like `@.`
+but does not materialise, returning a `Broadcasted` object for later use.
+Beware that mutation of arguments will affect the result,
+and that if it is used in two places, work will be done twice.
+"""
+macro lazy(ex)
+  bc = esc(Broadcast.__dot__(ex))
+  :($lazy.($bc))
 end
 
 function lazy end
@@ -106,6 +111,9 @@ Broadcast.broadcasted(::typeof(lazy), x) = Lazy(x)
 struct Lazy{T}; bc::T; end
 Broadcast.materialize(x::Lazy) = Broadcast.instantiate(x.bc)
 
+onevalue(λ::T, x::AbstractArray{T}) where T = map(_ -> λ, x)
+onevalue(λ, x::AbstractArray{T}) where T = onevalue(convert(float(T), λ), x)
+
 function Base.show(io::IO, ℓ::Leaf)  # show method is mostly to hide its long type!
   ioc = IOContext(io, :compact => true)
   print(ioc, "Leaf(", ℓ.rule, ", ")
diff --git a/src/rules.jl b/src/rules.jl
index f2e1ee5..71c138e 100644
--- a/src/rules.jl
+++ b/src/rules.jl
@@ -18,7 +18,7 @@ init(o::Descent, x::AbstractArray) = nothing
 function apply!(o::Descent, state, x, dx)
   η = convert(float(eltype(x)), o.eta)
   
-  return state, @.. dx * η
+  return state, @lazy dx * η  # @lazy creates a Broadcasted, will later fuse with x .= x .- dx
 end
 
 """
@@ -41,10 +41,10 @@ Momentum(η = 1f-2, ρ = 9f-1) = Momentum{typeof(η)}(η, ρ)
 init(o::Momentum, x::AbstractArray) = zero(x)
 
 function apply!(o::Momentum, state, x, dx)
-  η, ρ, v = o.eta, o.rho, state
-  v′ = @.. v = ρ * v - η * dx
+  η, ρ, mvel = o.eta, o.rho, state
+  @.. mvel = ρ * mvel + η * dx  # Macro @.. broadcasts into mvel if it can, else @. of rhs.
   
-  return v′, @.. -v′
+  return mvel, mvel
 end
 
 """
@@ -67,11 +67,12 @@ Nesterov(η = 1f-3, ρ = 9f-1) = Nesterov{typeof(η)}(η, ρ)
 init(o::Nesterov, x::AbstractArray) = zero(x)
 
 function apply!(o::Nesterov, state, x, dx)
-  η, ρ, v = o.eta, o.rho, state
-  d = @.. ρ^2 * v - (1+ρ) * η * dx
-  v′ = @.. v = ρ * v - η * dx
+  η, ρ, vel = o.eta, o.rho, state
+
+  newdx = @. - ρ^2 * vel + (1+ρ) * η * dx  # Cannot be lazy as this needs the old velocity
+  @.. vel = ρ * vel - η * dx
   
-  return v′, @.. -d
+  return vel, newdx
 end
 
 """
@@ -101,10 +102,11 @@ init(o::RMSProp, x::AbstractArray) = zero(x)
 
 function apply!(o::RMSProp, state, x, dx)
   η, ρ, ϵ, acc = o.eta, o.rho, o.epsilon, state
-  acc′ = @.. acc = ρ * acc + (1 - ρ) * dx^2
-  dx′ = @.. dx * (η / (sqrt(acc) + ϵ))
+
+  @.. acc = ρ * acc + (1 - ρ) * dx^2
+  dx′ = @lazy dx * (η / (sqrt(acc) + ϵ))
   
-  return acc′, dx′
+  return acc, dx′
 end
 
 """
@@ -129,15 +131,15 @@ ADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = ADAM{typeof(η)}(
 
 init(o::ADAM, x::AbstractArray) = (zero(x), zero(x), o.beta)
 
-function apply!(o::ADAM{T}, state, x, dx) where T
+function apply!(o::ADAM, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
   mt, vt, βt = state
 
-  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])) + ϵ) * η
+  @.. mt = β[1] * mt + (1 - β[1]) * dx
+  @.. vt = β[2] * vt + (1 - β[2]) * dx ^ 2
+  dx′ = @lazy mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ) * η
 
-  return (mt′, vt′, βt .* β), dx′
+  return (mt, vt, βt .* β), dx′
 end
 
 """
@@ -168,17 +170,17 @@ function apply!(o::RADAM, state, x, dx)
 
   mt, vt, βt, t = state
 
-  mt′ = @.. mt = β[1] * mt + (1 - β[1]) * dx
-  vt′ = @.. vt = β[2] * vt + (1 - β[2]) * dx^2
+  @.. mt = β[1] * mt + (1 - β[1]) * dx
+  @.. 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′ = @lazy mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ) * η * r
   else
-    dx′ = @.. mt / (1 - βt[1]) * η
+    dx′ = @lazy mt / (1 - βt[1]) * η
   end
 
-  return (mt′, vt′, βt .* β, t + 1), dx′
+  return (mt, vt, βt .* β, t + 1), dx′
 end
 
 """
@@ -205,14 +207,13 @@ init(o::AdaMax, x::AbstractArray) = (zero(x), zero(x), o.beta)
 
 function apply!(o::AdaMax, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
-
   mt, ut, βt = state
 
-  mt′ = @.. mt = β[1] * mt + (1 - β[1]) * dx
-  ut′ = @.. ut = max(β[2] * ut, abs(dx))
-  dx′ = @.. (η/(1 - βt[1])) * mt/(ut + ϵ)
+  @.. mt = β[1] * mt + (1 - β[1]) * dx
+  @.. ut = max(β[2] * ut, abs(dx))
+  dx′ = @lazy (η/(1 - βt[1])) * mt/(ut + ϵ)
 
-  return (mt′, ut′, βt .* β), dx′
+  return (mt, ut, βt .* β), dx′
 end
 
 """
@@ -240,16 +241,15 @@ init(o::OADAM, x::AbstractArray) = (zero(x), zero(x), o.beta, zero(x))
 
 function apply!(o::OADAM, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
+  mt, vt, βt, term = state
 
-  mt, vt, βt, dx_ = state
+  @.. mt = β[1] * mt + (1 - β[1]) * dx
+  @.. vt = β[2] * vt + (1 - β[2]) * dx^2
+  prev = copy(term)
+  @.. term = η * mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ)
+  dx′ = @lazy 2 * term - prev
 
-  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 .* β, term), dx′
 end
 
 """
@@ -271,16 +271,16 @@ struct ADAGrad{T}
 end
 ADAGrad(η = 1f-1, ϵ = eps(typeof(η))) = ADAGrad{typeof(η)}(η, ϵ)
 
-init(o::ADAGrad, x::AbstractArray) = fill!(similar(x), o.epsilon)
+init(o::ADAGrad, x::AbstractArray) = onevalue(o.epsilon, x)
 
 function apply!(o::ADAGrad, state, x, dx)
   η, ϵ = o.eta, o.epsilon
   acc = state
 
-  acc′ = @.. acc = acc + dx^2
-  dx′ = @.. dx * η / (sqrt(acc) + ϵ)
+  @.. acc = acc + dx^2
+  dx′ = @lazy dx * η / (sqrt(acc) + ϵ)
 
-  return acc′, dx′
+  return acc, dx′
 end
 
 """
@@ -307,13 +307,12 @@ function apply!(o::ADADelta, state, x, dx)
   ρ, ϵ = o.rho, o.epsilon
   acc, Δacc = state
 
-  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 = ρ * Δacc + (1 - ρ) * dx^2
+  @.. acc = ρ * acc + (1 - ρ) * dx^2
+  # DON'T remove epsilon from numerator or even out of the square roots!
+  dx′ = @. dx * sqrt(Δacc + ϵ) / sqrt(acc + ϵ)  # Cannot be lazy as this needs the old Δacc
+  @.. Δacc = ρ * Δacc + (1 - ρ) * dx′^2
   
-  return (acc′, Δacc′), dx′
+  return (acc, Δacc), dx′
 end
 
 """
@@ -338,19 +337,18 @@ end
 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))
+  (onevalue(o.epsilon, x), onevalue(o.epsilon, x), onevalue(o.epsilon, x))
 
 function apply!(o::AMSGrad, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
-
   mt, vt, v̂t = state
 
-  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) + ϵ)
+  @.. mt = β[1] * mt + (1 - β[1]) * dx
+  @.. vt = β[2] * vt + (1 - β[2]) * dx ^ 2
+  @.. v̂t = max(v̂t, vt)
+  dx′ = @lazy η * mt / (sqrt(v̂t) + ϵ)
 
-  return (mt′, vt′, v̂t′), dx′
+  return (mt, vt, v̂t), dx′
 end
 
 """
@@ -381,12 +379,12 @@ function apply!(o::NADAM, state, x, dx)
 
   mt, vt, βt = state
 
-  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])) / 
+  @.. mt = β[1] * mt + (1 - β[1]) * dx
+  @.. vt = β[2] * vt + (1 - β[2]) * dx^2
+  dx′ = @lazy (β[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
 
 """
@@ -405,7 +403,7 @@ weight decay regularization.
                          (no need to change default)
 """
 ADAMW(η = 1f-3, β = (9f-1, 9.99f-1), γ = 0, ϵ = eps(typeof(η))) =
-  OptimiserChain(ADAM{typeof(η)}(η, β, ϵ), WeightDecay(γ))
+  OptimiserChain(ADAM{typeof(η)}(η, β, ϵ), WeightDecay{typeof(η)}(γ))
 
 """
     AdaBelief(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))
@@ -434,11 +432,11 @@ function apply!(o::AdaBelief, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
   mt, st = state
 
-  mt′ = @.. mt = β[1] * mt + (1 - β[1]) * dx
-  st′ = @.. st = β[2] * st + (1 - β[2]) * (dx - mt)^2
-  dx′ = @.. η * mt / (sqrt(st) + ϵ)
+  @.. mt = β[1] * mt + (1 - β[1]) * dx
+  @.. st = β[2] * st + (1 - β[2]) * (dx - mt)^2
+  dx′ = @lazy η * mt / (sqrt(st) + ϵ)
   
-  return (mt′, st′), dx′
+  return (mt, st), dx′
 end
 
 """
@@ -457,7 +455,7 @@ WeightDecay() = WeightDecay(5f-4)
 init(o::WeightDecay, x::AbstractArray) = nothing
 
 function apply!(o::WeightDecay, state, x, dx)
-  dx′ = @.. dx + o.wd * x
+  dx′ = @lazy dx + o.wd * x
 
   return state, dx′
 end
@@ -478,7 +476,7 @@ init(o::ClipGrad, x::AbstractArray) = nothing
 
 function apply!(o::ClipGrad, state, x, dx)
   δ = convert(float(eltype(x)), o.delta)
-  dx′ = @.. clamp(dx, -δ, δ)
+  dx′ = @lazy clamp(dx, -δ, δ)
 
   return state, dx′
 end
@@ -510,7 +508,7 @@ function apply!(o::ClipNorm, state, x, dx)
   end
   λ = min(o.omega / nrm, 1)
 
-  return state, @.. dx * λ
+  return state, @lazy dx * λ
 end
 
 """


Test patch:
diff --git a/test/rules.jl b/test/rules.jl
index c950aeb..1d43b5a 100644
--- a/test/rules.jl
+++ b/test/rules.jl
@@ -5,7 +5,7 @@ using LinearAlgebra, Statistics, Test, Random
 Random.seed!(1)
 
 RULES = [
-  # All the rules at default settings
+  # All the rules at default settings:
   Descent(), ADAM(), Momentum(), Nesterov(), RMSProp(),
   ADAGrad(), AdaMax(), ADADelta(), AMSGrad(), NADAM(),
   ADAMW(), RADAM(), OADAM(), AdaBelief(),
@@ -13,6 +13,7 @@ RULES = [
   OptimiserChain(WeightDecay(), ADAM(0.001)),
   OptimiserChain(ClipNorm(), ADAM(0.001)),
   OptimiserChain(ClipGrad(0.5), Momentum()),
+  OptimiserChain(WeightDecay(), OADAM(), ClipGrad(1)),
 ]
 
 name(o) = typeof(o).name.name
@@ -63,13 +64,7 @@ end
       gs = gradient(x -> loss(x, w′), w)
       st, w = Optimisers.update(st, w, gs...)
     end
-    lw = loss(w, w′)
-    if o isa ADADelta
-      @show name(o) loss(w, w′)
-      @test_broken lw < 0.001
-    else
-      @test lw < 0.001
-    end
+    @test loss(w, w′) < 0.001
   end
 end
 
@@ -90,7 +85,7 @@ end
       g = gradient(m -> s_loss(m, x, y), model)[1]
       state, model = Optimisers.update!(state, model, g)
     end
-    if o isa Union{Descent, RMSProp, ADAGrad, ADADelta, NADAM}
+    if o isa Descent
       @show name(o) s_loss(model, x, y)
       @test_broken s_loss(model, x, y) < 1
     else
@@ -118,7 +113,7 @@ end
     end
     @test upstatic[1] isa SVector
 
-    # With ordinary Array gradient, what happens?
+    # With ordinary Array gradient, what happens? Not so important!
     upstatic2 = Optimisers.update(Optimisers.setup(o, mstatic), mstatic, marray[1:2])[2]
     # @test map(eltype, upstatic2) == types[1:2]  # same information
     if upstatic2[1] isa SVector
@@ -142,7 +137,7 @@ end
     @test xfill.a != ones(2,2)
     @test xfill.b != ones(2,2)
 
-    bc = Optimisers.@.. 1 + log([2 3; 4 5]) / 6
+    bc = Optimisers.@lazy 1 + log([2 3; 4 5]) / 6
     _, xbc = Optimisers.update(s, x, (a = bc, b = bc))
     @test xbc.a != ones(2,2)
     @test xbc.b != ones(2,2)
@@ -154,3 +149,20 @@ end
   end
 end
 
+@testset verbose=true "mutation check" begin
+  # If @lazy captures a matrix which is later mutated, the results won't agree here:
+  @testset "$(name(o))" for o in RULES
+    model = Float64.(rand(Int8, 8))
+    s_model = SVector{8}(model)
+    grads = [Float64.(rand(Int8, 8)) for t in 1:13]
+    s_grads = [SVector{8}(x) for x in grads]
+    state = Optimisers.setup(o, model)
+    s_state = Optimisers.setup(o, s_model)
+    for t in 1:13
+      state, model = Optimisers.update!(state, model, grads[t])
+      s_state, s_model = Optimisers.update!(s_state, s_model, s_grads[t])
+    end
+    @test model == s_model
+  end
+end
+
diff --git a/test/runtests.jl b/test/runtests.jl
index 7981034..ef94b10 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,7 +1,7 @@
 using Optimisers
 using ChainRulesCore, Functors, StaticArrays, Zygote
 using LinearAlgebra, Statistics, Test, Random
-using Optimisers: @..
+using Optimisers: @.., @lazy
 
 Random.seed!(1)
 
@@ -116,10 +116,10 @@ Optimisers.trainable(x::TwoThirds) = (a = x.a,)
       @test Optimisers.update!(s, m, g...)[2] isa Foo
     end
 
-    @testset "broadcasting macro" begin
+    @testset "broadcasting macros" begin
       x = [1.0, 2.0]; y = [3,4]; z = [5,6]
-      @test (@.. x + y * z) isa Broadcast.Broadcasted
-      bc = @.. x + y * z
+      @test (@lazy x + y * z) isa Broadcast.Broadcasted
+      bc = @lazy x + y * z
       @test (y .+ 2 .* bc) == [35,56]
 
       @test (@.. x = y * z) isa Array
@@ -127,6 +127,8 @@ Optimisers.trainable(x::TwoThirds) = (a = x.a,)
       r = 1.0:2.0  # immutable
       @test (@.. r = y * z) isa Array
       @test (@.. r = y * z) == y .* z
+      @.. r = y * z
+      @test r == y .* z  # attaches name r to result
     end
 
     @testset "tied weights" begin