Merge pull request #68 from SciML/dataarg

Don't use DiffResults in Flux optimiser dispatch with FiniteDiff

Author: Vaibhavdixit02

Project.toml

@@ -1,7 +1,7 @@
 name = "GalacticOptim"
 uuid = "a75be94c-b780-496d-a8a9-0878b188d577"
 authors = ["Vaibhavdixit02 <[email protected]>"]
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/function.jl

@@ -212,13 +212,13 @@ function instantiate_function(f, x, adtype::AutoFiniteDiff, p, num_cons = 0)
     _f = (θ, args...) -> first(f.f(θ, p, args...))
 
     if f.grad === nothing
-        grad = (res, θ, args...) -> FiniteDiff.finite_difference_gradient!(res,x ->_f(x, args...), θ, FiniteDiff.GradientCache(res, x, adtype.fdtype))
+        grad = (res, θ, args...) -> FiniteDiff.finite_difference_gradient!(res, x ->_f(x, args...), θ, FiniteDiff.GradientCache(res, x, adtype.fdtype))
     else
         grad = f.grad
     end
 
     if f.hess === nothing
-        hess = (res, θ, args...) -> FiniteDiff.finite_difference_hessian!(res,x ->_f(x, args...), θ, FiniteDiff.HessianCache(x, adtype.fdhtype))
+        hess = (res, θ, args...) -> FiniteDiff.finite_difference_hessian!(res, x ->_f(x, args...), θ, FiniteDiff.HessianCache(x, adtype.fdhtype))
     else
         hess = f.hess
     end

src/solve.jl

@@ -94,7 +94,7 @@ function __solve(prob::OptimizationProblem, opt, _data = DEFAULT_DATA;cb = (args
 
 	@withprogress progress name="Training" begin
 	  for (i,d) in enumerate(data)
-		gs = DiffResults.GradientResult(θ)
+		gs = prob.f.adtype isa AutoFiniteDiff ? Array{Number}(undef,length(θ)) : DiffResults.GradientResult(θ)
 		f.grad(gs, θ, d...) 
 		x = f.f(θ, prob.p, d...)
 		cb_call = cb(θ, x...)
@@ -105,7 +105,7 @@ function __solve(prob::OptimizationProblem, opt, _data = DEFAULT_DATA;cb = (args
 		end
 		msg = @sprintf("loss: %.3g", x[1])
 		progress && ProgressLogging.@logprogress msg i/maxiters
-		update!(opt, ps, DiffResults.gradient(gs))
+		update!(opt, ps, prob.f.adtype isa AutoFiniteDiff ? gs : DiffResults.gradient(gs))
 
 		if save_best
 		  if first(x) < first(min_err)  #found a better solution

test/ADtests.jl

@@ -107,3 +107,6 @@ sol = solve(prob, Newton())
 
 sol = solve(prob, Optim.KrylovTrustRegion())
 @test sol.minimum < l1 #the loss doesn't go below 5e-1 here
+
+sol = solve(prob, ADAM(0.1))
+@test 10*sol.minimum < l1