Repository: FluxML/Optimisers.jl (julia)

Issue:
Adam optimizer can produce NaNs with Float16 due to small epsilon
Dear All,

I have found a bug in Adam's optimiser when used with Float16. 
An MWE woud look like this
```julia
using Optimisers

x = Float16[0.579, -0.729, 0.5493]
δx = Float16[-0.001497, 0.0001875, -0.013176]

os = Optimisers.setup(Optimisers.Adam(Float16(1e-4)), x);
os, x = Optimisers.update(os, x, δx)
```
where after update `x = Float16[Inf, -Inf, 0.5493]`. This happens with `Optimisers 0.3.2`. I believe the problem is caused on this line https://github.com/FluxML/Optimisers.jl/blob/1908a1cd599f656b15304a9722328bf9b2eed360/src/rules.jl#L221 because default `ϵ` is in Float64 and it is too small.
When I initialize Adam with different epsilon as `Optimisers.Adam(1e-4, (0.9, 0.999),  eps(Float16))`, it works well.
A possible solution would be to make sure that the epsilon has at least some value with respect to the given type of Float. For example  change `T(o.epsilon)` to `max(T(o.epsilon), eps(T))` on line 216.

I can prepare a pull-request with this test.

Pull request:
Fix #167 adding `ϵ = max(ϵ, eps(T(0)))`. 

Also fix #189

Patch:
diff --git a/.gitignore b/.gitignore
index 952f7cea..763dd6f5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,3 +2,4 @@ Manifest.toml
 .vscode/
 docs/build/
 .DS_Store
+/test.jl
\ No newline at end of file
diff --git a/src/Optimisers.jl b/src/Optimisers.jl
index 2e115c40..99fc162f 100644
--- a/src/Optimisers.jl
+++ b/src/Optimisers.jl
@@ -25,6 +25,8 @@ export Descent, Adam, Momentum, Nesterov, Rprop, RMSProp,
        WeightDecay, SignDecay, ClipGrad, ClipNorm, OptimiserChain, Lion,
        AccumGrad
 
+VERSION >= v"1.11.0-DEV.469" && eval(Meta.parse("public apply!, init, setup, update, update!"))
+
 ###
 ### one-array functions
 ###
diff --git a/src/rules.jl b/src/rules.jl
index bc9c099b..b4fbd2a4 100644
--- a/src/rules.jl
+++ b/src/rules.jl
@@ -130,7 +130,7 @@ RMSProp(; eta = 0.001, rho = 0.9, epsilon = 1e-8, kw...) = RMSProp(eta, rho, eps
 init(o::RMSProp, x::AbstractArray) = (zero(x), o.centred ? zero(x) : false)
 
 function apply!(o::RMSProp, state, x::AbstractArray{T}, dx) where T
-  η, ρ, ϵ = T(o.eta), T(o.rho), T(o.epsilon)
+  η, ρ, ϵ = T(o.eta), T(o.rho), _eps(T, o.epsilon)
   quad, lin = state
 
   @.. quad = ρ * quad + (1 - ρ) * abs2(dx)
@@ -216,7 +216,7 @@ end
 init(o::Adam, x::AbstractArray{T}) where T = (zero(x), zero(x), T.(o.beta))
 
 function apply!(o::Adam, state, x::AbstractArray{T}, dx) where T
-  η, β, ϵ = T(o.eta), T.(o.beta), T(o.epsilon)
+  η, β, ϵ = T(o.eta), T.(o.beta), _eps(T, o.epsilon)
   mt, vt, βt = state
 
   @.. mt = β[1] * mt + (1 - β[1]) * dx
@@ -279,7 +279,7 @@ end
 init(o::RAdam, x::AbstractArray{T}) where T = (zero(x), zero(x), T.(o.beta), 1)
 
 function apply!(o::RAdam, state, x::AbstractArray{T}, dx) where T
-  η, β, ϵ = T(o.eta), T.(o.beta), T(o.epsilon)
+  η, β, ϵ = T(o.eta), T.(o.beta), _eps(T, o.epsilon)
   ρ∞ = 2/(1-β[2]) - 1 |> real
 
   mt, vt, βt, t = state
@@ -320,7 +320,7 @@ end
 init(o::AdaMax, x::AbstractArray{T}) where T = (zero(x), zero(x), T.(o.beta))
 
 function apply!(o::AdaMax, state, x::AbstractArray{T}, dx) where T
-  η, β, ϵ = T(o.eta), T.(o.beta), T(o.epsilon)
+  η, β, ϵ = T(o.eta), T.(o.beta), _eps(T, o.epsilon)
   mt, ut, βt = state
 
   @.. mt = β[1] * mt + (1 - β[1]) * dx
@@ -354,7 +354,7 @@ end
 init(o::OAdam, x::AbstractArray{T}) where T = (zero(x), zero(x), T.(o.beta), zero(x))
 
 function apply!(o::OAdam, state, x::AbstractArray{T}, dx) where T
-  η, β, ϵ = T(o.eta), T.(o.beta), T(o.epsilon)
+  η, β, ϵ = T(o.eta), T.(o.beta), _eps(T, o.epsilon)
   mt, vt, βt, term = state
 
   @.. mt = β[1] * mt + (1 - β[1]) * dx
@@ -388,7 +388,7 @@ end
 init(o::AdaGrad, x::AbstractArray) = onevalue(o.epsilon, x)
 
 function apply!(o::AdaGrad, state, x::AbstractArray{T}, dx) where T
-  η, ϵ = T(o.eta), T(o.epsilon)
+  η, ϵ = T(o.eta), _eps(T, o.epsilon)
   acc = state
 
   @.. acc = acc + abs2(dx)
@@ -418,7 +418,7 @@ end
 init(o::AdaDelta, x::AbstractArray) = (zero(x), zero(x))
 
 function apply!(o::AdaDelta, state, x::AbstractArray{T}, dx) where T
-  ρ, ϵ = T(o.rho), T(o.epsilon)
+  ρ, ϵ = T(o.rho), _eps(T, o.epsilon)
   acc, Δacc = state
 
   @.. acc = ρ * acc + (1 - ρ) * abs2(dx)
@@ -454,7 +454,7 @@ init(o::AMSGrad, x::AbstractArray) =
   (onevalue(o.epsilon, x), onevalue(o.epsilon, x), onevalue(o.epsilon, x))
 
 function apply!(o::AMSGrad, state, x::AbstractArray{T}, dx) where T
-  η, β, ϵ = T(o.eta), T.(o.beta), T(o.epsilon)
+  η, β, ϵ = T(o.eta), T.(o.beta), _eps(T, o.epsilon)
   mt, vt, v̂t = state
 
   @.. mt = β[1] * mt + (1 - β[1]) * dx
@@ -489,8 +489,7 @@ end
 init(o::NAdam, x::AbstractArray{T}) where T = (zero(x), zero(x), T.(o.beta))
 
 function apply!(o::NAdam, state, x::AbstractArray{T}, dx) where T
-  η, β, ϵ = T(o.eta), T.(o.beta), T(o.epsilon)
-
+  η, β, ϵ = T(o.eta), T.(o.beta), _eps(T, o.epsilon)
   mt, vt, βt = state
 
   @.. mt = β[1] * mt + (1 - β[1]) * dx
@@ -548,7 +547,7 @@ end
 init(o::AdaBelief, x::AbstractArray{T}) where T = (zero(x), zero(x), T.(o.beta))
 
 function apply!(o::AdaBelief, state, x::AbstractArray{T}, dx) where T
-  η, β, ϵ = T(o.eta), T.(o.beta), T(o.epsilon)
+  η, β, ϵ = T(o.eta), T.(o.beta), _eps(T, o.epsilon)
   mt, st, βt = state
 
   @.. mt = β[1] * mt + (1 - β[1]) * dx
diff --git a/src/utils.jl b/src/utils.jl
index 12a19dd0..8f667465 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -14,3 +14,9 @@ foreachvalue(f, x::Dict, ys...) = foreach(pairs(x)) do (k, v)
 end
 
 ofeltype(x, y) = convert(float(eltype(x)), y)
+
+_eps(T::Type{<:AbstractFloat}, e) = T(e)
+# catch complex and integers
+_eps(T::Type{<:Number}, e) = _eps(real(float(T)), e) 
+# avoid small e being rounded to zero
+_eps(T::Type{Float16}, e) = e == 0 ? T(0) : max(T(1e-7), T(e))


Test patch:
diff --git a/test/rules.jl b/test/rules.jl
index 52a3580e..9068fa13 100644
--- a/test/rules.jl
+++ b/test/rules.jl
@@ -267,3 +267,13 @@ end
   tree, x4 = Optimisers.update(tree, x3, g4)
   @test x4 ≈ x3
 end
+
+@testset "Float16 epsilon" begin
+  # issue https://github.com/FluxML/Optimisers.jl/issues/167
+  x = Float16[0.579, -0.729, 0.5493]
+  δx = Float16[-0.001497, 0.0001875, -0.013176]
+
+  os = Optimisers.setup(Adam(1e-4), x);
+  os, x = Optimisers.update(os, x, δx)
+  @test x ≈ Float16[1.835, -0.886, 0.5493] rtol=1e-3
+end