Some type annotations inspired by some JET errors. (#1194)

pkofod · web-flow · commit 3e4f8cd7fe17 · 2025-11-02T19:55:15.000+01:00
* Some type annotations inspired by some JET errors.

Not all these were related to JET errors, but since they appered because of unclear types in certain signatures I fiigured I'd add a few.

* Update gradient_descent.jl

* Options-&gt;NamedTuple

* Update golden_section.jl

* Update types.jl

* Update types.jl
diff --git a/src/multivariate/solvers/constrained/ipnewton/ipnewton.jl b/src/multivariate/solvers/constrained/ipnewton/ipnewton.jl
@@ -121,7 +121,7 @@ end
 
 function initial_state(
     method::IPNewton,
-    options,
+    options::Options,
     d::TwiceDifferentiable,
     constraints::TwiceDifferentiableConstraints,
     initial_x::AbstractArray{T},
@@ -270,7 +270,7 @@ function update_state!(
     constraints::TwiceDifferentiableConstraints,
     state::IPNewtonState{T},
     method::IPNewton,
-    options,
+    options::Options,
 ) where {T}
     state.f_x_previous, state.L_previous = state.f_x, state.L
     bstate, bstep, bounds = state.bstate, state.bstep, constraints.bounds
@@ -331,7 +331,7 @@ end
 function solve_step!(
     state::IPNewtonState,
     constraints,
-    options,
+    options::Options,
     show_linesearch::Bool = false,
 )
     x, s, μ, bounds = state.x, state.s, state.μ, constraints.bounds
diff --git a/src/multivariate/solvers/first_order/accelerated_gradient_descent.jl b/src/multivariate/solvers/first_order/accelerated_gradient_descent.jl
@@ -88,10 +88,10 @@ end
 function trace!(
     tr,
     d,
-    state,
-    iteration,
+    state::AcceleratedGradientDescentState,
+    iteration::Integer,
     method::AcceleratedGradientDescent,
-    options,
+    options::Options,
     curr_time = time(),
 )
     common_trace!(tr, d, state, iteration, method, options, curr_time)
diff --git a/src/multivariate/solvers/first_order/adam.jl b/src/multivariate/solvers/first_order/adam.jl
@@ -67,7 +67,7 @@ function _get_init_params(method::Adam)
     method.α(1), method.β₁, method.β₂
 end
 
-function initial_state(method::Adam, options, d, initial_x::AbstractArray{T}) where {T}
+function initial_state(method::Adam, options::Options, d, initial_x::AbstractArray{T}) where {T}
     initial_x = copy(initial_x)
 
     value_gradient!!(d, initial_x)
@@ -122,6 +122,6 @@ function update_state!(d, state::AdamState{T}, method::Adam) where {T}
     false # break on linesearch error
 end
 
-function trace!(tr, d, state, iteration, method::Adam, options, curr_time = time())
+function trace!(tr, d, state::AdamState, iteration::Integer, method::Adam, options::Options, curr_time = time())
     common_trace!(tr, d, state, iteration, method, options, curr_time)
 end
diff --git a/src/multivariate/solvers/first_order/adamax.jl b/src/multivariate/solvers/first_order/adamax.jl
@@ -67,7 +67,7 @@ function _get_init_params(method::AdaMax)
     method.α(1), method.β₁, method.β₂
 end
 
-function initial_state(method::AdaMax, options, d, initial_x::AbstractArray{T}) where {T}
+function initial_state(method::AdaMax, options::Options, d, initial_x::AbstractArray{T}) where {T}
     initial_x = copy(initial_x)
 
     value_gradient!!(d, initial_x)
@@ -115,6 +115,6 @@ function update_state!(d, state::AdaMaxState{T}, method::AdaMax) where {T}
     false # break on linesearch error
 end
 
-function trace!(tr, d, state, iteration, method::AdaMax, options, curr_time = time())
+function trace!(tr, d, state::AdaMaxState, iteration::Integer, method::AdaMax, options::Options, curr_time = time())
     common_trace!(tr, d, state, iteration, method, options, curr_time)
 end
diff --git a/src/multivariate/solvers/first_order/bfgs.jl b/src/multivariate/solvers/first_order/bfgs.jl
@@ -88,7 +88,7 @@ function reset!(method, state::BFGSState, obj, x)
     end
 end
 
-function initial_state(method::BFGS, options, d, initial_x::AbstractArray{T}) where {T}
+function initial_state(method::BFGS, options::Options, d, initial_x::AbstractArray{T}) where {T}
     n = length(initial_x)
     initial_x = copy(initial_x)
     retract!(method.manifold, initial_x)
@@ -186,7 +186,7 @@ function update_h!(d, state, method::BFGS)
     end
 end
 
-function trace!(tr, d, state, iteration, method::BFGS, options, curr_time = time())
+function trace!(tr, d, state::BFGSState, iteration::Integer, method::BFGS, options::Options, curr_time = time())
     dt = Dict()
     dt["time"] = curr_time
     if options.extended_trace
diff --git a/src/multivariate/solvers/first_order/cg.jl b/src/multivariate/solvers/first_order/cg.jl
@@ -113,7 +113,7 @@ function reset!(cg::ConjugateGradient, cgs::ConjugateGradientState, obj, x)
     cgs.s .= .-cgs.pg
     cgs.f_x_previous = typeof(cgs.f_x_previous)(NaN)
 end
-function initial_state(method::ConjugateGradient, options, d, initial_x)
+function initial_state(method::ConjugateGradient, options::Options, d, initial_x)
     T = eltype(initial_x)
     initial_x = copy(initial_x)
     retract!(method.manifold, initial_x)
diff --git a/src/multivariate/solvers/first_order/gradient_descent.jl b/src/multivariate/solvers/first_order/gradient_descent.jl
@@ -97,10 +97,10 @@ end
 function trace!(
     tr,
     d,
-    state,
-    iteration,
+    state::GradientDescentState,
+    iteration::Integer,
     method::GradientDescent,
-    options,
+    options::Options,
     curr_time = time(),
 )
     common_trace!(tr, d, state, iteration, method, options, curr_time)
diff --git a/src/multivariate/solvers/first_order/l_bfgs.jl b/src/multivariate/solvers/first_order/l_bfgs.jl
@@ -152,14 +152,14 @@ mutable struct LBFGSState{Tx,Tdx,Tdg,T,G} <: AbstractOptimizerState
     s::Tx
     @add_linesearch_fields()
 end
-function reset!(method, state::LBFGSState, obj, x)
+function reset!(method::LBFGS, state::LBFGSState, obj, x::AbstractArray)
     retract!(method.manifold, x)
     value_gradient!(obj, x)
     project_tangent!(method.manifold, gradient(obj), x)
 
     state.pseudo_iteration = 0
 end
-function initial_state(method::LBFGS, options, d, initial_x)
+function initial_state(method::LBFGS, options::Options, d, initial_x::AbstractArray)
     T = real(eltype(initial_x))
     n = length(initial_x)
     initial_x = copy(initial_x)
@@ -228,7 +228,7 @@ function update_state!(d, state::LBFGSState, method::LBFGS)
 end
 
 
-function update_h!(d, state, method::LBFGS)
+function update_h!(d, state::LBFGSState, method::LBFGS)
     n = length(state.x)
     # Measure the change in the gradient
     state.dg .= gradient(d) .- state.g_previous
@@ -247,6 +247,6 @@ function update_h!(d, state, method::LBFGS)
     false
 end
 
-function trace!(tr, d, state, iteration, method::LBFGS, options, curr_time = time())
+function trace!(tr, d, state::LBFGSState, iteration::Integer, method::LBFGS, options::Options, curr_time = time())
     common_trace!(tr, d, state, iteration, method, options, curr_time)
 end
diff --git a/src/multivariate/solvers/first_order/momentum_gradient_descent.jl b/src/multivariate/solvers/first_order/momentum_gradient_descent.jl
@@ -28,7 +28,7 @@ mutable struct MomentumGradientDescentState{Tx,T} <: AbstractOptimizerState
     @add_linesearch_fields()
 end
 
-function initial_state(method::MomentumGradientDescent, options, d, initial_x)
+function initial_state(method::MomentumGradientDescent, options::Options, d, initial_x::AbstractArray)
     T = eltype(initial_x)
     initial_x = copy(initial_x)
     retract!(method.manifold, initial_x)
@@ -70,10 +70,10 @@ end
 function trace!(
     tr,
     d,
-    state,
-    iteration,
+    state::MomentumGradientDescentState,
+    iteration::Integer,
     method::MomentumGradientDescent,
-    options,
+    options::Options,
     curr_time = time(),
 )
     common_trace!(tr, d, state, iteration, method, options, curr_time)
diff --git a/src/multivariate/solvers/first_order/ngmres.jl b/src/multivariate/solvers/first_order/ngmres.jl
@@ -223,7 +223,7 @@ end
 const ngmres_oaccel_warned = Ref{Bool}(false)
 function initial_state(
     method::AbstractNGMRES,
-    options,
+    options::Options,
     d,
     initial_x::AbstractArray{eTx},
 ) where {eTx}
@@ -452,10 +452,10 @@ end
 function trace!(
     tr,
     d,
-    state,
-    iteration,
+    state::NGMRESState,
+    iteration::Integer,
     method::AbstractNGMRES,
-    options,
+    options::Options,
     curr_time = time(),
 )
     dt = Dict()
diff --git a/src/multivariate/solvers/second_order/krylov_trust_region.jl b/src/multivariate/solvers/second_order/krylov_trust_region.jl
@@ -36,7 +36,7 @@ mutable struct KrylovTrustRegionState{T} <: AbstractOptimizerState
     cg_iters::Int
 end
 
-function initial_state(method::KrylovTrustRegion, options, d, initial_x::Array{T}) where {T}
+function initial_state(method::KrylovTrustRegion, options::Options, d, initial_x::Array{T}) where {T}
     n = length(initial_x)
     # Maintain current gradient in gr
     @assert(method.max_radius > 0)
@@ -68,10 +68,10 @@ end
 function trace!(
     tr,
     d,
-    state,
-    iteration,
+    state::KrylovTrustRegionState,
+    iteration::Integer,
     method::KrylovTrustRegion,
-    options,
+    options::Options,
     curr_time = time(),
 )
     dt = Dict()
diff --git a/src/multivariate/solvers/second_order/newton.jl b/src/multivariate/solvers/second_order/newton.jl
@@ -87,7 +87,7 @@ function update_state!(d, state::NewtonState, method::Newton)
     return !lssuccess # break on linesearch error
 end
 
-function trace!(tr, d, state, iteration, method::Newton, options, curr_time = time())
+function trace!(tr, d, state::NewtonState, iteration::Integer, method::Newton, options::Options, curr_time = time())
     dt = Dict()
     dt["time"] = curr_time
     if options.extended_trace
diff --git a/src/multivariate/solvers/second_order/newton_trust_region.jl b/src/multivariate/solvers/second_order/newton_trust_region.jl
@@ -420,10 +420,10 @@ end
 function trace!(
     tr,
     d,
-    state,
-    iteration,
+    state::NewtonTrustRegionState,
+    iteration::Integer,
     method::NewtonTrustRegion,
-    options,
+    options::Options,
     curr_time = time(),
 )
     dt = Dict()
diff --git a/src/types.jl b/src/types.jl
@@ -189,17 +189,8 @@ function Options(;
 end
 
 function Options(o::Options; kws...)
-    fields = fieldnames(typeof(o))
-    foreach(keys(kws)) do kw
-        if kw ∉ fields
-            error(lazy"`$kw` is not a valid keyword argument of `Options(opts; kws...)`")
-        end
-    end
-    newargs = ntuple(Val(nfields(o))) do i
-        field = fields[i]
-        get(kws, field, getfield(o, field))
-    end
-    return Options(newargs...)
+    o_nt = (; (name => getproperty(o, name) for name in propertynames(o))...)
+    return Options(; o_nt..., kws...)
 end
 
 _show_helper(output, k, v) = output * "$k = $v, "
diff --git a/src/univariate/solvers/brent.jl b/src/univariate/solvers/brent.jl
@@ -216,7 +216,7 @@ function optimize(
 end
 
 
-function trace!(tr, d, state, iteration, method::Brent, options, curr_time = time())
+function trace!(tr, d::Nothing, state::NamedTuple, iteration::Integer, method::Brent, options::NamedTuple, curr_time = time())
     dt = Dict()
     dt["time"] = curr_time
     dt["minimizer"] = state.new_minimizer
diff --git a/src/univariate/solvers/golden_section.jl b/src/univariate/solvers/golden_section.jl
@@ -154,7 +154,7 @@ function optimize(
 end
 
 
-function trace!(tr, d, state, iteration, method::GoldenSection, options, curr_time = time())
+function trace!(tr, d, state, iteration::Integer, method::GoldenSection, options::NamedTuple, curr_time = time())
     dt = Dict()
     dt["time"] = curr_time
     dt["minimizer"] = state.new_minimizer
diff --git a/src/utilities/trace.jl b/src/utilities/trace.jl
@@ -4,9 +4,9 @@ function common_trace!(
     tr,
     d,
     state,
-    iteration,
+    iteration::Integer,
     method::FirstOrderOptimizer,
-    options,
+    options::Options,
     curr_time = time(),
 )
     dt = Dict()
diff --git a/test/general/types.jl b/test/general/types.jl
@@ -53,6 +53,6 @@ import Compat.String
     @test opts2.f_calls_limit == 100 && opts2.f_calls_limit != opts1.f_calls_limit
     @test opts2.f_abstol == opts1.f_abstol == 1e-6
     @test opts2.iterations == opts1.iterations == 1000
-    @test_throws "not a valid keyword argument" Optim.Options(opts1; invalid_keyword=1)
+    @test_throws MethodError Optim.Options(opts1; invalid_keyword=1)
     @test Optim.Options(opts1) == opts1
 end
