Make manually passed derivatives be 3 argument functions

Author: Vaibhavdixit02

src/function/finitediff.jl

@@ -59,7 +59,7 @@ function instantiate_function(f, x, adtype::AutoFiniteDiff, p, num_cons = 0)
                                                                                    args...),
                                                                            θ, gradcache)
     else
-        grad = f.grad
+        grad = (G, θ, args...) -> f.grad(G, θ, p, args...)
     end
 
     if f.hess === nothing
@@ -71,7 +71,7 @@ function instantiate_function(f, x, adtype::AutoFiniteDiff, p, num_cons = 0)
                                                                           updatecache(hesscache,
                                                                                       θ))
     else
-        hess = f.hess
+        hess = (H, θ, args...) -> f.hess(H, θ, p, args...)
     end
 
     if f.hv === nothing
@@ -102,7 +102,7 @@ function instantiate_function(f, x, adtype::AutoFiniteDiff, p, num_cons = 0)
             FiniteDiff.finite_difference_jacobian!(J, cons, θ, jaccache)
         end
     else
-        cons_j = f.cons_j
+        cons_j = (J, θ) -> f.cons_j(J, θ, p)
     end
 
     if cons !== nothing && f.cons_h === nothing
@@ -120,13 +120,11 @@ function instantiate_function(f, x, adtype::AutoFiniteDiff, p, num_cons = 0)
             end
         end
     else
-        cons_h = f.cons_h
+        cons_h = (res, θ) -> f.cons_h(res, θ, p)
     end
 
-    return OptimizationFunction{true}(f, adtype; grad = grad, hess = hess, hv = hv,
-                                      cons = cons, cons_j = cons_j, cons_h = cons_h,
-                                      cons_jac_colorvec = cons_jac_colorvec,
-                                      hess_prototype = nothing,
-                                      cons_jac_prototype = f.cons_jac_prototype,
-                                      cons_hess_prototype = nothing)
+    return OptimizationFunction{true}(f, adtype; grad=grad, hess=hess, hv=hv,
+        cons=cons, cons_j=cons_j, cons_h=cons_h,
+        cons_jac_colorvec = cons_jac_colorvec,
+        hess_prototype=f.hess_prototype, cons_jac_prototype=f.cons_jac_prototype, cons_hess_prototype=f.cons_hess_prototype)
 end

src/function/forwarddiff.jl

@@ -51,15 +51,15 @@ function instantiate_function(f::OptimizationFunction{true}, x,
         grad = (res, θ, args...) -> ForwardDiff.gradient!(res, x -> _f(x, args...), θ,
                                                           gradcfg, Val{false}())
     else
-        grad = f.grad
+        grad = (G, θ, args...) -> f.grad(G, θ, p, args...)
     end
 
     if f.hess === nothing
         hesscfg = ForwardDiff.HessianConfig(_f, x, ForwardDiff.Chunk{chunksize}())
         hess = (res, θ, args...) -> ForwardDiff.hessian!(res, x -> _f(x, args...), θ,
                                                          hesscfg, Val{false}())
     else
-        hess = f.hess
+        hess = (H, θ, args...) -> f.hess(H, θ, p, args...)
     end
 
     if f.hv === nothing
@@ -85,7 +85,7 @@ function instantiate_function(f::OptimizationFunction{true}, x,
             ForwardDiff.jacobian!(J, cons_oop, θ, cjconfig)
         end
     else
-        cons_j = f.cons_j
+        cons_j = (J, θ) -> f.cons_j(J, θ, p)
     end
 
     if cons !== nothing && f.cons_h === nothing
@@ -99,12 +99,10 @@ function instantiate_function(f::OptimizationFunction{true}, x,
             end
         end
     else
-        cons_h = f.cons_h
+        cons_h = (res, θ) -> f.cons_h(res, θ, p)
     end
 
-    return OptimizationFunction{true}(f.f, adtype; grad = grad, hess = hess, hv = hv,
-                                      cons = cons, cons_j = cons_j, cons_h = cons_h,
-                                      hess_prototype = nothing,
-                                      cons_jac_prototype = f.cons_jac_prototype,
-                                      cons_hess_prototype = f.cons_hess_prototype)
+    return OptimizationFunction{true}(f.f, adtype; grad=grad, hess=hess, hv=hv,
+        cons=cons, cons_j=cons_j, cons_h=cons_h,
+        hess_prototype=f.hess_prototype, cons_jac_prototype=f.cons_jac_prototype, cons_hess_prototype=f.cons_hess_prototype)
 end

src/function/mtk.jl

@@ -15,15 +15,15 @@ function instantiate_function(f, x, adtype::AutoModelingToolkit, p, num_cons = 0
         grad_oop, grad_iip = ModelingToolkit.generate_gradient(sys, expression = Val{false})
         grad(J, u) = (grad_iip(J, u, p); J)
     else
-        grad = f.grad
+        grad = (G, θ, args...) -> f.grad(G, θ, p, args...)
     end
 
     if f.hess === nothing
         hess_oop, hess_iip = ModelingToolkit.generate_hessian(sys, expression = Val{false},
                                                               sparse = adtype.obj_sparse)
         hess(H, u) = (hess_iip(H, u, p); H)
     else
-        hess = f.hess
+        hess = (H, θ, args...) -> f.hess(H, θ, p, args...)
     end
 
     if f.hv === nothing
@@ -69,7 +69,7 @@ function instantiate_function(f, x, adtype::AutoModelingToolkit, p, num_cons = 0
             jac_iip(J, θ, p)
         end
     else
-        cons_j = f.cons_j
+        cons_j = (J, θ) -> f.cons_j(J, θ, p)
     end
 
     if f.cons !== nothing && f.cons_h === nothing
@@ -82,7 +82,7 @@ function instantiate_function(f, x, adtype::AutoModelingToolkit, p, num_cons = 0
             cons_hess_iip(res, θ, p)
         end
     else
-        cons_h = f.cons_h
+        cons_h = (res, θ) -> f.cons_h(res, θ, p)
     end
 
     if adtype.obj_sparse

src/function/reversediff.jl

@@ -53,7 +53,7 @@ function instantiate_function(f, x, adtype::AutoReverseDiff, p = SciMLBase.NullP
         grad = (res, θ, args...) -> ReverseDiff.gradient!(res, x -> _f(x, args...), θ,
                                                           ReverseDiff.GradientConfig(θ))
     else
-        grad = f.grad
+        grad = (G, θ, args...) -> f.grad(G, θ, p, args...)
     end
 
     if f.hess === nothing
@@ -70,7 +70,7 @@ function instantiate_function(f, x, adtype::AutoReverseDiff, p = SciMLBase.NullP
             end
         end
     else
-        hess = f.hess
+        hess = (H, θ, args...) -> f.hess(H, θ, p, args...)
     end
 
     if f.hv === nothing
@@ -84,9 +84,7 @@ function instantiate_function(f, x, adtype::AutoReverseDiff, p = SciMLBase.NullP
         hv = f.hv
     end
 
-    return OptimizationFunction{false}(f, adtype; grad = grad, hess = hess, hv = hv,
-                                       cons = nothing, cons_j = nothing, cons_h = nothing,
-                                       hess_prototype = nothing,
-                                       cons_jac_prototype = nothing,
-                                       cons_hess_prototype = nothing)
+    return OptimizationFunction{false}(f, adtype; grad=grad, hess=hess, hv=hv,
+        cons=nothing, cons_j=nothing, cons_h=nothing,
+        hess_prototype=f.hess_prototype, cons_jac_prototype=nothing, cons_hess_prototype=nothing)
 end

src/function/tracker.jl

@@ -36,13 +36,13 @@ function instantiate_function(f, x, adtype::AutoTracker, p, num_cons = 0)
                                     res .= Tracker.data(Tracker.gradient(x -> _f(x, args...),
                                                                          θ)[1])
     else
-        grad = f.grad
+        grad = (G, θ, args...) -> f.grad(G, θ, p, args...)
     end
 
     if f.hess === nothing
         hess = (res, θ, args...) -> error("Hessian based methods not supported with Tracker backend, pass in the `hess` kwarg")
     else
-        hess = f.hess
+        hess = (H, θ, args...) -> f.hess(H, θ, p, args...)
     end
 
     if f.hv === nothing
@@ -51,9 +51,8 @@ function instantiate_function(f, x, adtype::AutoTracker, p, num_cons = 0)
         hv = f.hv
     end
 
-    return OptimizationFunction{false}(f, adtype; grad = grad, hess = hess, hv = hv,
-                                       cons = nothing, cons_j = nothing, cons_h = nothing,
-                                       hess_prototype = nothing,
-                                       cons_jac_prototype = nothing,
-                                       cons_hess_prototype = nothing)
+
+    return OptimizationFunction{false}(f, adtype; grad=grad, hess=hess, hv=hv,
+        cons=nothing, cons_j=nothing, cons_h=nothing,
+        hess_prototype=f.hess_prototype, cons_jac_prototype=nothing, cons_hess_prototype=nothing)
 end

src/function/zygote.jl

@@ -36,7 +36,7 @@ function instantiate_function(f, x, adtype::AutoZygote, p, num_cons = 0)
                                                                           θ)[1]) :
                                     res .= Zygote.gradient(x -> _f(x, args...), θ)[1]
     else
-        grad = f.grad
+        grad = (G, θ, args...) -> f.grad(G, θ, p, args...)
     end
 
     if f.hess === nothing
@@ -52,7 +52,7 @@ function instantiate_function(f, x, adtype::AutoZygote, p, num_cons = 0)
             end
         end
     else
-        hess = f.hess
+        hess = (H, θ, args...) -> f.hess(H, θ, p, args...)
     end
 
     if f.hv === nothing
@@ -66,9 +66,7 @@ function instantiate_function(f, x, adtype::AutoZygote, p, num_cons = 0)
         hv = f.hv
     end
 
-    return OptimizationFunction{false}(f, adtype; grad = grad, hess = hess, hv = hv,
-                                       cons = nothing, cons_j = nothing, cons_h = nothing,
-                                       hess_prototype = nothing,
-                                       cons_jac_prototype = nothing,
-                                       cons_hess_prototype = nothing)
+    return OptimizationFunction{false}(f, adtype; grad=grad, hess=hess, hv=hv,
+        cons=nothing, cons_j=nothing, cons_h=nothing,
+        hess_prototype=f.hess_prototype, cons_jac_prototype=nothing, cons_hess_prototype=nothing)
 end