Repository: FluxML/Optimisers.jl (julia)

Issue:
Type Constraints in the Rule Structs
I was working on getting Optimisers to work with reactant, and it (mostly) does, but one of the current issues is that `eta` is forced to be Float64 in some of the structs.

But consider the following IR:

```julia
julia> pss = (; a = (rand(3) |> Reactant.to_rarray))
(a = ConcreteRArray{Float64, 1}([0.023549009580651203, 0.10813549621409191, 0.7874517465499301]),)

julia> st_opt = @allowscalar Optimisers.setup(opt, pss)
(a = Leaf(Adam(0.001, (0.9, 0.999), 1.0e-8), (ConcreteRArray{Float64, 1}([0.0, 0.0, 0.0]), ConcreteRArray{Float64, 1}([0.0, 0.0, 0.0]), (0.9, 0.999))),)

julia> @code_hlo Optimisers.update(st_opt, pss, pss)
module {
  func.func @main(%arg0: tensor<3xf64>, %arg1: tensor<3xf64>, %arg2: tensor<3xf64>) -> (tensor<3xf64>, tensor<3xf64>, tensor<3xf64>) {
    %cst = stablehlo.constant dense<1.000000e-03> : tensor<3xf64>
    %cst_0 = stablehlo.constant dense<1.000000e-08> : tensor<3xf64>
    %cst_1 = stablehlo.constant dense<0.0010000000000000009> : tensor<3xf64>
    %cst_2 = stablehlo.constant dense<0.99899999999999999> : tensor<3xf64>
    %cst_3 = stablehlo.constant dense<0.099999999999999978> : tensor<3xf64>
    %cst_4 = stablehlo.constant dense<9.000000e-01> : tensor<3xf64>
    %0 = stablehlo.multiply %cst_4, %arg0 : tensor<3xf64>
    %1 = stablehlo.multiply %cst_3, %arg2 : tensor<3xf64>
    %2 = stablehlo.add %0, %1 : tensor<3xf64>
    %3 = stablehlo.multiply %cst_2, %arg1 : tensor<3xf64>
    %4 = stablehlo.abs %arg2 : tensor<3xf64>
    %5 = stablehlo.multiply %4, %4 : tensor<3xf64>
    %6 = stablehlo.multiply %cst_1, %5 : tensor<3xf64>
    %7 = stablehlo.add %3, %6 : tensor<3xf64>
    %8 = stablehlo.divide %2, %cst_3 : tensor<3xf64>
    %9 = stablehlo.divide %7, %cst_1 : tensor<3xf64>
    %10 = stablehlo.sqrt %9 : tensor<3xf64>
    %11 = stablehlo.add %10, %cst_0 : tensor<3xf64>
    %12 = stablehlo.divide %8, %11 : tensor<3xf64>
    %13 = stablehlo.multiply %12, %cst : tensor<3xf64>
    %14 = stablehlo.subtract %arg2, %13 : tensor<3xf64>
    return %2, %7, %14 : tensor<3xf64>, tensor<3xf64>, tensor<3xf64>
  }
}
```

While this looks correct, if you take a closer look:

```
%cst_1 = stablehlo.constant dense<0.0010000000000000009> : tensor<3xf64>
%cst_2 = stablehlo.constant dense<0.99899999999999999> : tensor<3xf64>
%cst_3 = stablehlo.constant dense<0.099999999999999978> : tensor<3xf64>
```

The learning rate and the other parameters get embedded into the IR as constants. So even if we do an `adjust` of the learning rate, it will still be using the old learning rate

Pull request description:
Fix #205

Patch:
diff --git a/Project.toml b/Project.toml
index 1c6de35e..e5e89f07 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,10 +1,11 @@
 name = "Optimisers"
 uuid = "3bd65402-5787-11e9-1adc-39752487f4e2"
-version = "0.4.4"
 authors = ["Mike J Innes <mike.j.innes@gmail.com>"]
+version = "0.4.4"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -21,16 +22,8 @@ OptimisersEnzymeCoreExt = "EnzymeCore"
 [compat]
 Adapt = "4"
 ChainRulesCore = "1"
+ConstructionBase = "1.5.8"
 EnzymeCore = "0.8.5"
 Functors = "0.4.9, 0.5"
 Statistics = "1"
-Zygote = "0.6.40, 0.7.1"
 julia = "1.10"
-
-[extras]
-StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
-
-[targets]
-test = ["Test", "EnzymeCore", "StaticArrays", "Zygote"]
diff --git a/src/Optimisers.jl b/src/Optimisers.jl
index 99fc162f..6456f20c 100644
--- a/src/Optimisers.jl
+++ b/src/Optimisers.jl
@@ -4,6 +4,7 @@ using Functors: functor, fmap, fmap_with_path,
                 KeyPath, haskeypath, getkeypath,
                 isleaf, @functor, fmapstructure, children, AbstractWalk
 using LinearAlgebra
+using ConstructionBase: ConstructionBase
 
 include("interface.jl")
 export AbstractRule
diff --git a/src/adjust.jl b/src/adjust.jl
index c0ce9366..214d70f7 100644
--- a/src/adjust.jl
+++ b/src/adjust.jl
@@ -142,5 +142,5 @@ function _adjust(r::T, nt::NamedTuple) where T <: AbstractRule
   vals = map(fs) do field
     get(nt, field, getfield(r, field))
   end
-  T(vals...)  # relies on having the default constructor
+  return ConstructionBase.constructorof(T)(vals...)
 end
diff --git a/src/interface.jl b/src/interface.jl
index e44dec15..68ae0b42 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -239,39 +239,69 @@ Helper macro for defining rules with default values.
 The types of the literal values are used in the `struct`,
 like this:
 ```julia
-struct Rule
-  eta::Float64
-  beta::Tuple{Float64, Float64}
+struct Rule{T1, T2}
+  eta::T1
+  beta::T2
+
   Rule(eta, beta = (0.7, 0.8)) = eta < 0 ? error() : new(eta, beta)
   Rule(; eta = 0.1, beta = (0.7, 0.8)) = Rule(eta, beta)
 end
+
 ```
 Any field called `eta` is assumed to be a learning rate, and cannot be negative.
 """
 macro def(expr)
   Meta.isexpr(expr, :struct) || throw("@def must act on a struct definition")
   lines = expr.args[3].args
-  names, vals = [], []
+  names, default_vals = [], []
+  default_types, type_params = [], []
   for i in eachindex(lines)
     lines[i] isa Symbol && throw("@def requires a default for every field")
     Meta.isexpr(lines[i], :(=)) || continue
     name, val = lines[i].args
     push!(names, name)
-    push!(vals, val)
-    lines[i] = :($name::$typeof($val))
+    push!(default_vals, val)
+    push!(default_types, _def_typeof(val))
+    # @show name, val, typeof(val)
+    type = Symbol("T$name")
+    push!(type_params, type)
+    lines[i] = :($name::$type)
   end
   rule = Meta.isexpr(expr.args[2], :<:) ? expr.args[2].args[1] : expr.args[2]
-  params = [Expr(:kw, nv...) for nv in zip(names,vals)]
-  check = :eta in names ? :(eta < 0 && throw(DomainError(eta, "the learning rate cannot be negative"))) : nothing
+  expr.args[2] = Expr(:(<:), Expr(:curly, rule, type_params...), :AbstractRule)
+  params = [Expr(:kw, nv...) for nv in zip(names, default_vals)]
+  check_sign_eta = :eta in names ? :(eta < 0 && throw(DomainError(eta, "the learning rate cannot be negative"))) : nothing
   # Positional-argument method, has defaults for all but the first arg:
-  inner = :(function $rule($(names[1]), $(params[2:end]...))
-    $check
-    new($(names...))
+  positional = :(function $rule($(names[1]), $(params[2:end]...))
+    vars = $(_def_check_type).([$(names...)],[$(default_types...)])
+    $check_sign_eta
+    return new{typeof.(vars)...}(vars...)
   end)
   # Keyword-argument method. (Made an inner constructor only to allow
   # resulting structs to be @doc... cannot if macro returns a block.)
   kwmethod = :($rule(; $(params...)) = $rule($(names...)))
-  push!(lines, inner, kwmethod)
-  esc(expr)
+  push!(lines, positional, kwmethod)
+  # return esc(expr)
+# end
+
+  return quote 
+          Base.@__doc__ $expr
+
+          function Base.show(io::IO, r::$rule)
+            pairs = ["$n=$(getfield(r, n))" for n in [($names...)]]
+            print(io, $rule,"(", join(pairs, ", "), ")")
+          end
+        end |> esc
 end
 
+_def_typeof(val) = typeof(val)
+_def_typeof(val::Expr) = typeof(eval(val))
+
+_def_check_type(x, T) = x
+_def_check_type(x, T::Type{<:Number}) = throw(ArgumentError("$x is not a number"))
+_def_check_type(x::Number, T::Type{<:Number}) = x
+_def_check_type(x::Number, T::Type{<:AbstractFloat}) = float(x)
+_def_check_type(x::Complex, T::Type{<:AbstractFloat}) = throw(ArgumentError("cannot convert complex to real"))
+_def_check_type(x, T::Type{<:AbstractFloat}) = throw(ArgumentError("$x is not a number"))
+_def_check_type(x::Tuple, T::Type{<:Tuple}) = _def_check_type.(x, (T.parameters...,))
+_def_check_type(x, T::Type{<:Tuple}) = throw(ArgumentError("$x is not a tuple"))
diff --git a/src/rules.jl b/src/rules.jl
index 0cd8d30c..fa2ad1ef 100644
--- a/src/rules.jl
+++ b/src/rules.jl
@@ -114,16 +114,16 @@ gradients by an estimate their variance, instead of their second moment.
 - Keyword `centred` (or `centered`): Indicates whether to use centred variant
                                      of the algorithm.
 """
-struct RMSProp <: AbstractRule
-  eta::Float64
-  rho::Float64
-  epsilon::Float64
+struct RMSProp{Teta,Trho,Teps} <: AbstractRule
+  eta::Teta
+  rho::Trho
+  epsilon::Teps
   centred::Bool
 end
 
 function RMSProp(η, ρ = 0.9, ϵ = 1e-8; centred::Bool = false, centered::Bool = false)
   η < 0 && throw(DomainError(η, "the learning rate cannot be negative"))
-  RMSProp(η, ρ, ϵ, centred | centered)
+  return RMSProp(float(η), float(ρ), float(ϵ), centred | centered)
 end
 RMSProp(; eta = 0.001, rho = 0.9, epsilon = 1e-8, kw...) = RMSProp(eta, rho, epsilon; kw...)
 
@@ -148,14 +148,12 @@ function adjust(r::RMSProp; kw...)
 end
 
 function Base.show(io::IO, o::RMSProp)
-  print(io, "RMSProp(")
-  join(io, [o.eta, o.rho, o.epsilon], ", ")
-  print(io, "; centred = ", o.centred, ")")
+  print(io, "RMSProp(eta=$(o.eta), rho=$(o.rho), epsilon=$(o.epsilon), centred=$(o.centred)")
 end
 
 
 """
-    Rprop(η = 1f-3, ℓ = (5f-1, 1.2f0), Γ = (1f-6, 50f0))
+    Rprop(η = 1e-3, ℓ = (0.5, 1.2), Γ = (1e-6, 50.0))
     Rprop(; [eta, ell, gamma])
 
 Optimizer using the
@@ -171,9 +169,9 @@ learning algorithm that depends only on the sign of the gradient.
 - Step sizes (`Γ::Tuple == gamma`): Mminimal and maximal allowed step sizes.
 """
 @def struct Rprop <: AbstractRule
-    eta =  1f-3
-    ell = (5f-1, 1.2f0)
-    gamma = (1f-6, 50f0)
+    eta =  1e-3
+    ell = (0.5, 1.2)
+    gamma = (1e-6, 50.0)
 end
 
 init(o::Rprop, x::AbstractArray) = (zero(x), onevalue(o.eta, x))
@@ -528,17 +526,17 @@ Implemented as an [`OptimiserChain`](@ref) of [`Adam`](@ref) and [`WeightDecay`]
     The previous rule, which is closer to the original paper, can be obtained by setting `AdamW(..., couple=false)`.
     See [this issue](https://github.com/FluxML/Flux.jl/issues/2433) for more details.
 """
-struct AdamW <: AbstractRule
-  eta::Float64
-  beta::Tuple{Float64, Float64}
-  lambda::Float64
-  epsilon::Float64
+struct AdamW{Teta,Tbeta<:Tuple,Tlambda,Teps} <: AbstractRule
+  eta::Teta
+  beta::Tbeta
+  lambda::Tlambda
+  epsilon::Teps
   couple::Bool
 end
 
 function AdamW(η, β = (0.9, 0.999), λ = 0.0, ϵ = 1e-8; couple::Bool = true)
   η < 0 && throw(DomainError(η, "the learning rate cannot be negative"))
-  AdamW(η, β, λ, ϵ, couple)
+  return AdamW(float(η), β, float(λ), float(ϵ), couple)
 end
 
 AdamW(; eta = 0.001, beta = (0.9, 0.999), lambda= 0.0,  epsilon = 1e-8, kw...) =
@@ -565,6 +563,11 @@ function apply!(o::AdamW, state, x::AbstractArray{T}, dx) where T
   return (mt, vt, βt .* β), dx′′
 end
 
+function Base.show(io::IO, o::AdamW)
+  print(io, "AdamW(eta=$(o.eta), beta=$(o.beta), lambda=$(o.lambda), epsilon=$(o.epsilon), couple=$(o.couple))")
+end
+
+
 """
     AdaBelief(η = 0.001, β = (0.9, 0.999), ϵ = 1e-16)
     AdaBelief(; [eta, beta, epsilon])
@@ -704,12 +707,12 @@ Typically composed with other rules using [`OptimiserChain`](@ref).
 
 See also [`ClipGrad`](@ref).
 """
-struct ClipNorm <: AbstractRule
-  omega::Float64
-  p::Float64
+struct ClipNorm{To,Tp} <: AbstractRule
+  omega::To
+  p::Tp
   throw::Bool
 end
-ClipNorm(ω = 10, p = 2; throw::Bool = true) = ClipNorm(ω, p, throw)
+ClipNorm(ω = 10, p = 2; throw::Bool = true) = ClipNorm(float(ω), float(p), throw)
 
 init(o::ClipNorm, x::AbstractArray) = nothing
 


Test patch:
diff --git a/test/Project.toml b/test/Project.toml
new file mode 100644
index 00000000..b3c4b20a
--- /dev/null
+++ b/test/Project.toml
@@ -0,0 +1,11 @@
+[deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
diff --git a/test/destructure.jl b/test/destructure.jl
index 232a9001..b348507e 100644
--- a/test/destructure.jl
+++ b/test/destructure.jl
@@ -98,37 +98,12 @@ end
       sum(gradient(m -> sum(destructure(m)[1])^3, (v, [4,5,6.0]))[1][1])
     end[1] == [378, 378, 378]
 
-    VERSION >= v"1.10" && @test gradient([1,2,3.0]) do v
+    @test gradient([1,2,3.0]) do v
       sum(abs2, gradient(m -> sum(abs2, destructure(m)[1]), (x = v, y = sin, z = [4,5,6.0]))[1][1])
     end[1] ≈ [8,16,24]
     # Zygote error in (::typeof(∂(canonicalize)))(Δ::NamedTuple{(:backing,), Tuple{NamedTuple{(:x, :y, :z)
     # Diffractor error in perform_optic_transform
   end
-  
-  false && @testset "using Yota" begin
-    @test Yota_gradient(m -> destructure(m)[1][1], m1)[1] == [1,0,0]
-    @test Yota_gradient(m -> destructure(m)[1][2], m2)[1] == ([0,1,0], [0,0,0])
-    @test Yota_gradient(m -> destructure(m)[1][3], (m1, m1))[1] == ([0,0,1], nothing)
-    @test Yota_gradient(m -> destructure(m)[1][1], m3)[1] == (x = [1,0,0], y = nothing, z = [0,0,0])
-    @test Yota_gradient(m -> destructure(m)[1][2], m4)[1] == (x = [0,1,0], y = nothing, z = [0,0,0])
-
-    g5 = Yota_gradient(m -> destructure(m)[1][3], m5)[1]
-    @test g5.a[1].x == [0,0,1]
-    @test g5.a[2] === nothing
-
-    g6 = Yota_gradient(m -> imag(destructure(m)[1][4]), m6)[1]
-    @test g6.a == [0,0,0]
-    @test g6.a isa Vector{Float64}
-    @test g6.b == [0+im]
-
-    g8 = Yota_gradient(m -> sum(abs2, destructure(m)[1]), m8)[1]
-    @test g8[1].x == [2,4,6]
-    @test g8[2].b.x == [8]
-    @test g8[3] == [[10.0]]
-
-    g9 = Yota_gradient(m -> sum(sqrt, destructure(m)[1]), m9)[1]
-    @test g9.c === nothing
-  end
 end
 
 @testset "gradient of rebuild" begin
@@ -174,36 +149,6 @@ end
     # Not fixed by this:
     # Zygote.@adjoint Tangent{T,B}(x::NamedTuple) where {T,B<:NamedTuple} = Tangent{T,B}(x), dx -> (dx,)
   end
-  
-  false && @testset "using Yota" begin
-    re1 = destructure(m1)[2]
-    @test Yota_gradient(x -> re1(x)[1], rand(3))[1] == [1,0,0]
-    re2 = destructure(m2)[2]
-    @test Yota_gradient(x -> re2(x)[1][2], rand(6))[1] == [0,1,0,0,0,0]
-    re3 = destructure(m3)[2]
-    @test Yota_gradient(x -> re3(x).x[3], rand(6))[1] == [0,0,1,0,0,0]
-    @test Yota_gradient(x -> re3(x).z[1], rand(6))[1] == [0,0,0,1,0,0]
-
-    re4 = destructure(m4)[2]
-    @test Yota_gradient(x -> re4(x).x[1], rand(6))[1] == [1,0,0,0,0,0]
-    @test Yota_gradient(x -> re4(x).y[2], rand(6))[1] == [0,1,0,0,0,0]
-    @test Yota_gradient(rand(6)) do x
-      m = re4(x)
-      m.x[1] + 2*m.y[2] + 3*m.z[3]
-    end[1] == [1,2,0, 0,0,3]
-
-    re7 = destructure(m7)[2]
-    @test Yota_gradient(x -> re7(x).a[2][3], rand(3))[1] == [0,0,1]
-    @test Yota_gradient(x -> re7(x).b[2][2], rand(3))[1] == [0,0,0]
-    @test Yota_gradient(x -> re7(x).c[2][1], rand(3))[1] == [0,0,0]
-
-    v8, re8 = destructure(m8)
-    @test Yota_gradient(x -> sum(abs2, re8(x)[1].y), v8)[1] == [2,4,6,0,0]
-    @test Yota_gradient(x -> only(sum(re8(x)[3]))^2, v8)[1] == [0,0,0,0,10]
-
-    re9 = destructure(m9)[2]
-    @test Yota_gradient(x -> sum(abs2, re9(x).c[1]), 1:7)[1] == [0,0,0, 8,10,12,14]
-  end
 end
 
 @testset "Flux issue 1826" begin
diff --git a/test/interface.jl b/test/interface.jl
new file mode 100644
index 00000000..6878851d
--- /dev/null
+++ b/test/interface.jl
@@ -0,0 +1,30 @@
+@testset "@def" begin 
+    Optimisers.@def struct DummyRule
+        eta = 1.0
+        b1 = 1
+        b2 = 2.5f0
+        c = (1.0, 2.0)
+    end
+
+    # no args
+    r = DummyRule()
+    @test r.eta === 1.0
+    @test r.b1 === 1
+    @test r.b2 === 2.5f0
+    @test r.c === (1.0, 2.0)
+
+    # some positional args    
+    r = DummyRule(2, 2, 4.5)
+    @test r.eta === 2.0 # int promoted to float
+    @test r.b1 === 2
+    @test r.b2 === 4.5  # Float64 not converted to Float32
+    @test r.c === (1.0, 2.0)
+
+    # errors due to type mismatch
+    @test_throws ArgumentError DummyRule(0.1+im)
+    @test_throws ArgumentError DummyRule(b2=0.1+im)
+    @test_throws ArgumentError DummyRule(0.1, :beta)
+    @test_throws ArgumentError DummyRule(c=0.9)
+    @test_throws ArgumentError DummyRule(c=(0.9, 0.99im))
+    @test_throws ArgumentError Optimisers.adjust(DummyRule(), c=0.9)
+end
diff --git a/test/rules.jl b/test/rules.jl
index 499902ca..4194d093 100644
--- a/test/rules.jl
+++ b/test/rules.jl
@@ -230,21 +230,6 @@ end
   end
 end
 
-false && @testset "using Yota" begin
-  @testset "$(name(o))" for o in RULES
-    w′ = (abc = (α = rand(3, 3), β = rand(3, 3), γ = rand(3)), d = (δ = rand(3), ε = eps))
-    w = (abc = (α = 5rand(3, 3), β = rand(3, 3), γ = rand(3)), d = (δ = rand(3), ε = eps))
-    st = Optimisers.setup(o, w)
-    loss(x, y) = mean((x.abc.α .* x.abc.β .- y.abc.α .* y.abc.β) .^ 2)  # does not use γ, δ, ε
-    @test loss(w, w′) > 0.5
-    for i = 1:10^4
-      _, (_, g, _) = Yota.grad(loss, w, w′)
-      st, w = Optimisers.update(st, w, g)
-    end
-    @test loss(w, w′) < 0.001
-  end
-end
-
 @testset "AccumGrad" begin
   x0 = rand(5)
   x = copy(x0)
diff --git a/test/runtests.jl b/test/runtests.jl
index 956aa041..8cdcffce 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -38,15 +38,6 @@ function Optimisers.apply!(o::BiRule, state, x, dx, dx2)
   return state, dx
 end
 
-# if VERSION < v"1.9-"
-#   using Yota
-# end
-# # Make Yota's output look like Zygote's:
-
-# Yota_gradient(f, xs...) = map(y2z, Base.tail(Yota.grad(f, xs...)[2]))
-# y2z(::AbstractZero) = nothing  # we don't care about different flavours of zero
-# y2z(t::Tangent) = map(y2z, ChainRulesCore.backing(canonicalize(t)))  # namedtuples!
-# y2z(x) = x
 
 @testset verbose=true "Optimisers.jl" begin
   @testset verbose=true "Features" begin
@@ -566,4 +557,7 @@ end
   @testset verbose=true "Optimisation Rules" begin
     include("rules.jl")
   end
+  @testset verbose=true "interface" begin
+    include("interface.jl")
+  end
 end