Update optim fg, optimisers extensions and tests

Author: Vaibhavdixit02

lib/OptimizationOptimJL/src/OptimizationOptimJL.jl

@@ -26,6 +26,7 @@ function SciMLBase.requireshessian(opt::Union{
     true
 end
 SciMLBase.requiresgradient(opt::Optim.Fminbox) = true
+SciMLBase.allowsfg(opt::Union{Optim.AbstractOptimizer, Optim.ConstrainedOptimizer, Optim.Fminbox, Optim.SAMIN}) = true
 
 function __map_optimizer_args(cache::OptimizationCache,
         opt::Union{Optim.AbstractOptimizer, Optim.Fminbox,
@@ -142,11 +143,11 @@ function SciMLBase.__solve(cache::OptimizationCache{
         θ = metadata[cache.opt isa Optim.NelderMead ? "centroid" : "x"]
         opt_state = Optimization.OptimizationState(iter = trace.iteration,
             u = θ,
-            objective = x[1],
+            objective = trace.value,
             grad = get(metadata, "g(x)", nothing),
             hess = get(metadata, "h(x)", nothing),
             original = trace)
-        cb_call = cache.callback(opt_state, x...)
+        cb_call = cache.callback(opt_state, trace.value)
         if !(cb_call isa Bool)
             error("The callback should return a boolean `halt` for whether to stop the optimization process.")
         end
@@ -261,11 +262,11 @@ function SciMLBase.__solve(cache::OptimizationCache{
             metadata["x"]
         opt_state = Optimization.OptimizationState(iter = trace.iteration,
             u = θ,
-            objective = x[1],
+            objective = trace.value,
             grad = get(metadata, "g(x)", nothing),
             hess = get(metadata, "h(x)", nothing),
             original = trace)
-        cb_call = cache.callback(opt_state, x...)
+        cb_call = cache.callback(opt_state, trace.value)
         if !(cb_call isa Bool)
             error("The callback should return a boolean `halt` for whether to stop the optimization process.")
         end
@@ -277,15 +278,21 @@ function SciMLBase.__solve(cache::OptimizationCache{
         __x = first(x)
         return cache.sense === Optimization.MaxSense ? -__x : __x
     end
-    fg! = function (G, θ)
-        if G !== nothing
-            cache.f.grad(G, θ)
-            if cache.sense === Optimization.MaxSense
-                G .*= -one(eltype(G))
+
+    if cache.f.fg === nothing
+        fg! = function (G, θ)
+            if G !== nothing
+                cache.f.grad(G, θ)
+                if cache.sense === Optimization.MaxSense
+                    G .*= -one(eltype(G))
+                end
             end
+            return _loss(θ)
         end
-        return _loss(θ)
+    else
+        fg! = cache.f.fg
     end
+    
 
     gg = function (G, θ)
         cache.f.grad(G, θ)
@@ -344,9 +351,9 @@ function SciMLBase.__solve(cache::OptimizationCache{
             u = metadata["x"],
             grad = get(metadata, "g(x)", nothing),
             hess = get(metadata, "h(x)", nothing),
-            objective = x[1],
+            objective = trace.value,
             original = trace)
-        cb_call = cache.callback(opt_state, x...)
+        cb_call = cache.callback(opt_state, trace.value)
         if !(cb_call isa Bool)
             error("The callback should return a boolean `halt` for whether to stop the optimization process.")
         end
@@ -358,15 +365,21 @@ function SciMLBase.__solve(cache::OptimizationCache{
         __x = first(x)
         return cache.sense === Optimization.MaxSense ? -__x : __x
     end
-    fg! = function (G, θ)
-        if G !== nothing
-            cache.f.grad(G, θ)
-            if cache.sense === Optimization.MaxSense
-                G .*= -one(eltype(G))
+
+    if cache.f.fg === nothing
+        fg! = function (G, θ)
+            if G !== nothing
+                cache.f.grad(G, θ)
+                if cache.sense === Optimization.MaxSense
+                    G .*= -one(eltype(G))
+                end
             end
+            return _loss(θ)
         end
-        return _loss(θ)
+    else
+        fg! = cache.f.fg
     end
+
     gg = function (G, θ)
         cache.f.grad(G, θ)
         if cache.sense === Optimization.MaxSense
@@ -434,7 +447,7 @@ PrecompileTools.@compile_workload begin
     function obj_f(x, p)
         A = p[1]
         b = p[2]
-        return sum((A * x - b) .^ 2)
+        return sum((A * x .- b) .^ 2)
     end
 
     function solve_nonnegative_least_squares(A, b, solver)

lib/OptimizationOptimisers/Project.toml

@@ -10,17 +10,7 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 
-[weakdeps]
-MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
-MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
-
-[extensions]
-OptimizationOptimisersMLDataDevicesExt = "MLDataDevices"
-OptimizationOptimisersMLUtilsExt = "MLUtils"
-
 [compat]
-MLDataDevices = "1.1"
-MLUtils = "0.4.4"
 Optimisers = "0.2, 0.3"
 Optimization = "4"
 ProgressLogging = "0.1"

lib/OptimizationOptimisers/ext/OptimizationOptimisersMLDataDevicesExt.jl

@@ -2,7 +2,7 @@ module OptimizationOptimisers
 
 using Reexport, Printf, ProgressLogging
 @reexport using Optimisers, Optimization
-using Optimization.SciMLBase
+using Optimization.SciMLBase, Optimization.OptimizationBase
 
 SciMLBase.supports_opt_cache_interface(opt::AbstractRule) = true
 SciMLBase.requiresgradient(opt::AbstractRule) = true
@@ -16,8 +16,6 @@ function SciMLBase.__init(
         kwargs...)
 end
 
-isa_dataiterator(data) = false
-
 function SciMLBase.__solve(cache::OptimizationCache{
         F,
         RC,
@@ -59,7 +57,7 @@ function SciMLBase.__solve(cache::OptimizationCache{
         throw(ArgumentError("The number of epochs must be specified as the epochs or maxiters kwarg."))
     end
 
-    if isa_dataiterator(cache.p)
+    if OptimizationBase.isa_dataiterator(cache.p)
         data = cache.p
         dataiterate = true
     else

lib/OptimizationOptimisers/ext/OptimizationOptimisersMLUtilsExt.jl

@@ -31,11 +31,11 @@ end
 function loss_adjoint(p)
     prediction = predict_adjoint(p)
     loss = sum(abs2, x - 1 for x in prediction)
-    return loss, prediction
+    return loss
 end
 
 iter = 0
-callback = function (state, l, pred)
+callback = function (state, l)
     display(l)
 
     # using `remake` to re-create our `prob` with current parameters `p`
@@ -81,11 +81,11 @@ end
 function loss_neuralode(p)
     pred = predict_neuralode(p)
     loss = sum(abs2, ode_data .- pred)
-    return loss, pred
+    return loss
 end
 
 iter = 0
-callback = function (st, l, pred...)
+callback = function (st, l)
     global iter
     iter += 1
 

lib/OptimizationOptimisers/src/OptimizationOptimisers.jl

@@ -45,7 +45,7 @@ end
 function loss_adjoint(fullp, p)
     (batch, time_batch) = p
     pred = predict_adjoint(fullp, time_batch)
-    sum(abs2, batch .- pred), pred
+    sum(abs2, batch .- pred)
 end
 
 k = 10