From f870457bdc9b2ccfc90fe6e264cfa415ab93ae2f Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Mon, 26 Aug 2024 13:46:22 +0200
Subject: [PATCH] call cons_vjp if available

---
 lib/OptimizationMOI/Project.toml           |  2 ++
 lib/OptimizationMOI/src/OptimizationMOI.jl |  1 +
 lib/OptimizationMOI/src/nlp.jl             | 30 +++++++++++++---------
 3 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/lib/OptimizationMOI/Project.toml b/lib/OptimizationMOI/Project.toml
index 7b12d52b8..6195b9e14 100644
--- a/lib/OptimizationMOI/Project.toml
+++ b/lib/OptimizationMOI/Project.toml
@@ -4,6 +4,7 @@ authors = ["Vaibhav Dixit <vaibhavyashdixit@gmail.com> and contributors"]
 version = "0.4.2"
 
 [deps]
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
@@ -19,6 +20,7 @@ HiGHS = "1"
 Ipopt = "1"
 Ipopt_jll = "300.1400"
 Juniper = "0.9"
+LinearAlgebra = "1"
 MathOptInterface = "1"
 ModelingToolkit = "9"
 NLopt = "1"
diff --git a/lib/OptimizationMOI/src/OptimizationMOI.jl b/lib/OptimizationMOI/src/OptimizationMOI.jl
index 74e3bb4d9..72ae90165 100644
--- a/lib/OptimizationMOI/src/OptimizationMOI.jl
+++ b/lib/OptimizationMOI/src/OptimizationMOI.jl
@@ -11,6 +11,7 @@ import ModelingToolkit: parameters, unknowns, varmap_to_vars, mergedefaults, toe
 import ModelingToolkit
 const MTK = ModelingToolkit
 using Symbolics
+using LinearAlgebra
 
 const MOI = MathOptInterface
 
diff --git a/lib/OptimizationMOI/src/nlp.jl b/lib/OptimizationMOI/src/nlp.jl
index 152c5a7ba..934e14f5a 100644
--- a/lib/OptimizationMOI/src/nlp.jl
+++ b/lib/OptimizationMOI/src/nlp.jl
@@ -114,7 +114,7 @@ function MOIOptimizationNLPCache(prob::OptimizationProblem,
 
     num_cons = prob.ucons === nothing ? 0 : length(prob.ucons)
     f = Optimization.instantiate_function(prob.f, reinit_cache, prob.f.adtype, num_cons;
-    g = true, h = true, cons_j = true, lag_h = true)
+    g = true, h = false, cons_j = true, lag_h = true)
     T = eltype(prob.u0)
     n = length(prob.u0)
 
@@ -297,17 +297,23 @@ end
 #     return
 # end
 
-# function MOI.eval_constraint_jacobian_transpose_product(
-#     evaluator::Evaluator,
-#     y,
-#     x,
-#     w,
-# )
-#     start = time()
-#     MOI.eval_constraint_jacobian_transpose_product(evaluator.backend, y, x, w)
-#     evaluator.eval_constraint_jacobian_timer += time() - start
-#     return
-# end
+function MOI.eval_constraint_jacobian_transpose_product(
+    evaluator::MOIOptimizationNLPEvaluator,
+    y,
+    x,
+    w,
+)
+    if evaluator.f.cons_vjp !== nothing
+        evaluator.f.cons_vjp(y, w, x)
+
+    elseif evaluator.f.cons_j !== nothing
+        J = evaluator.J
+        evaluator.f.cons_j(J, x)
+        mul!(y, J', w)
+        return
+    end
+    error("Thou shalt provide the v'J of the constraint jacobian, not doing so is associated with great misfortune and also no ice cream for you.")
+end
 
 function MOI.hessian_lagrangian_structure(evaluator::MOIOptimizationNLPEvaluator)
     lagh = evaluator.f.lag_h !== nothing