Merge pull request #845 from aml5600/fix/vjp-return

MOI Constraint jacobian-vector fixes
SciML · Oct 30, 2024 · 5cb17b6 · 5cb17b6
2 parents d483b16 + ab79a1b
commit 5cb17b6
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 6 deletions.
diff --git a/lib/OptimizationMOI/src/nlp.jl b/lib/OptimizationMOI/src/nlp.jl
@@ -302,14 +302,14 @@ function MOI.eval_constraint_jacobian_product(
         evaluator::MOIOptimizationNLPEvaluator, y, x, w)
     if evaluator.f.cons_jvp !== nothing
         evaluator.f.cons_jvp(y, x, w)
-
     elseif evaluator.f.cons_j !== nothing
         J = evaluator.J
         evaluator.f.cons_j(J, x)
         mul!(y, J, w)
-        return
+    else
+        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
-    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.")
+    return nothing
 end
 
 function MOI.eval_constraint_jacobian_transpose_product(
@@ -320,14 +320,14 @@ function MOI.eval_constraint_jacobian_transpose_product(
 )
     if evaluator.f.cons_vjp !== nothing
         evaluator.f.cons_vjp(y, x, w)
-
     elseif evaluator.f.cons_j !== nothing
         J = evaluator.J
         evaluator.f.cons_j(J, x)
         mul!(y, J', w)
-        return
+    else
+        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
-    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.")
+    return nothing
 end
 
 function MOI.hessian_lagrangian_structure(evaluator::MOIOptimizationNLPEvaluator)

diff --git a/lib/OptimizationMOI/test/runtests.jl b/lib/OptimizationMOI/test/runtests.jl
@@ -2,6 +2,8 @@ using OptimizationMOI, Optimization, Ipopt, NLopt, Zygote, ModelingToolkit, Reve
 using AmplNLWriter, Ipopt_jll, Juniper, HiGHS
 using Test, SparseArrays
 
+import MathOptInterface
+
 function _test_sparse_derivatives_hs071(backend, optimizer)
     function objective(x, ::Any)
         return x[1] * x[4] * (x[1] + x[2] + x[3]) + x[3]
@@ -29,6 +31,32 @@ function _test_sparse_derivatives_hs071(backend, optimizer)
     return
 end
 
+@testset "Evaluator" begin
+    rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
+    x0 = zeros(2)
+    _p = [1.0, 100.0]
+    cons_circ = (res, x, p) -> res .= [x[1]^2 + x[2]^2]
+    optprob = OptimizationFunction(
+        rosenbrock, Optimization.AutoZygote();
+        cons = cons_circ)
+    prob = OptimizationProblem(optprob, x0, _p, ucons = [Inf], lcons = [0.0])
+    evaluator = init(prob, Ipopt.Optimizer()).evaluator
+
+    x = prob.u0
+    # vector-constraint jacobian product
+    @test (evaluator.f.cons_j !== nothing) || (evaluator.f.cons_jvp !== nothing)
+    y = zeros(1)
+    w = ones(2)
+    @test MathOptInterface.eval_constraint_jacobian_product(evaluator, y, x, w) === nothing
+
+    # constraint jacobian-vector product
+    @test (evaluator.f.cons_j !== nothing) || (evaluator.f.cons_vjp !== nothing)
+    y = zeros(2)
+    w = ones(1)
+    @test MathOptInterface.eval_constraint_jacobian_transpose_product(
+        evaluator, y, x, w) === nothing
+end
+
 @testset "NLP" begin
     rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
     x0 = zeros(2)