Make manually passed derivatives be 3 argument functions

Author: Vaibhavdixit02

src/function/finitediff.jl

@@ -54,14 +54,14 @@ function instantiate_function(f, x, adtype::AutoFiniteDiff, p, num_cons=0)
         gradcache = FiniteDiff.GradientCache(x, x, adtype.fdtype)
         grad = (res, θ, args...) -> FiniteDiff.finite_difference_gradient!(res, x -> _f(x, args...), θ, gradcache)
     else
-        grad = f.grad
+        grad = (G, θ, args...) -> f.grad(G, θ, p, args...)
     end
 
     if f.hess === nothing
         hesscache = FiniteDiff.HessianCache(x, adtype.fdhtype)
         hess = (res, θ, args...) -> FiniteDiff.finite_difference_hessian!(res, x -> _f(x, args...), θ, updatecache(hesscache, θ))
     else
-        hess = f.hess
+        hess = (H, θ, args...) -> f.hess(H, θ, p, args...)
     end
 
     if f.hv === nothing
@@ -89,7 +89,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
@@ -100,11 +100,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)
+        hess_prototype=f.hess_prototype, cons_jac_prototype=f.cons_jac_prototype, cons_hess_prototype=f.cons_hess_prototype)
 end

src/function/forwarddiff.jl

@@ -49,14 +49,14 @@ function instantiate_function(f::OptimizationFunction{true}, x, adtype::AutoForw
         gradcfg = ForwardDiff.GradientConfig(_f, x, ForwardDiff.Chunk{chunksize}())
         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
@@ -82,7 +82,7 @@ function instantiate_function(f::OptimizationFunction{true}, x, adtype::AutoForw
             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
@@ -94,10 +94,10 @@ function instantiate_function(f::OptimizationFunction{true}, x, adtype::AutoForw
             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)
+        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,14 +15,14 @@ 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
@@ -61,7 +61,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
@@ -70,7 +70,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

@@ -51,7 +51,7 @@ function instantiate_function(f, x, adtype::AutoReverseDiff, p=SciMLBase.NullPar
     if f.grad === nothing
         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
@@ -67,7 +67,7 @@ function instantiate_function(f, x, adtype::AutoReverseDiff, p=SciMLBase.NullPar
             end
         end
     else
-        hess = f.hess
+        hess = (H, θ, args...) -> f.hess(H, θ, p, args...)
     end
 
 
@@ -84,5 +84,5 @@ function instantiate_function(f, x, adtype::AutoReverseDiff, p=SciMLBase.NullPar
 
     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)
+        hess_prototype=f.hess_prototype, cons_jac_prototype=nothing, cons_hess_prototype=nothing)
 end

src/function/tracker.jl

@@ -30,13 +30,13 @@ function instantiate_function(f, x, adtype::AutoTracker, p, num_cons = 0)
     if f.grad === nothing
         grad = (res, θ, args...) -> res isa DiffResults.DiffResult ? DiffResults.gradient!(res, Tracker.data(Tracker.gradient(x -> _f(x, args...), θ)[1])) : 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
@@ -46,7 +46,7 @@ function instantiate_function(f, x, adtype::AutoTracker, p, num_cons = 0)
     end
 
 
-    return OptimizationFunction{false}(f, adtype; grad=grad, hess=hess, hv=hv, 
+    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)
+        hess_prototype=f.hess_prototype, cons_jac_prototype=nothing, cons_hess_prototype=nothing)
 end

src/function/zygote.jl

@@ -31,7 +31,7 @@ function instantiate_function(f, x, adtype::AutoZygote, p, num_cons = 0)
     if f.grad === nothing
         grad = (res, θ, args...) -> res isa DiffResults.DiffResult ? DiffResults.gradient!(res, Zygote.gradient(x -> _f(x, args...), θ)[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
@@ -47,7 +47,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
@@ -61,7 +61,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, 
+    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)
+        hess_prototype=f.hess_prototype, cons_jac_prototype=nothing, cons_hess_prototype=nothing)
 end