Interface for solver supporting diff natively (#356)

* Interface for solver supporting diff natively

* Fixes

* Fix format

* Claude improves coverage

* Fix format

* Add tests

* fix format

* Rename

* Use the attributes in inner models too

* Simplify

* Fix format

* Fix

* Restore keyword option for nonlinear

* Fix

* Fix

* Fixes

* Fix

* Fix

* Fix

* Fix

* Simpler

* Fix format

* Simplify handling of POI layer

* Fix format

Author: blegat

src/ConicProgram/ConicProgram.jl

@@ -253,7 +253,9 @@ function _gradient_cache(model::Model)
     return model.gradient_cache
 end
 
-function DiffOpt.forward_differentiate!(model::Model)
+MOI.supports(::Model, ::DiffOpt.ForwardDifferentiate) = true
+
+function MOI.set(model::Model, ::DiffOpt.ForwardDifferentiate, ::Nothing)
     model.diff_time = @elapsed begin
         gradient_cache = _gradient_cache(model)
         M = gradient_cache.M
@@ -333,7 +335,9 @@ function DiffOpt.forward_differentiate!(model::Model)
     # return -dx, -dy, -ds
 end
 
-function DiffOpt.reverse_differentiate!(model::Model)
+MOI.supports(::Model, ::DiffOpt.ReverseDifferentiate) = true
+
+function MOI.set(model::Model, ::DiffOpt.ReverseDifferentiate, ::Nothing)
     model.diff_time = @elapsed begin
         gradient_cache = _gradient_cache(model)
         M = gradient_cache.M

src/NonLinearProgram/NonLinearProgram.jl

@@ -512,6 +512,13 @@ function _cache_evaluator!(model::Model)
     return model.cache
 end
 
+MOI.supports(::Model, ::DiffOpt.ForwardDifferentiate) = true
+
+function MOI.set(model::Model, ::DiffOpt.ForwardDifferentiate, value)
+    kws = something(value, (;))
+    return DiffOpt.forward_differentiate!(model; kws...)
+end
+
 function DiffOpt.forward_differentiate!(model::Model; tol = 1e-6)
     model.diff_time = @elapsed begin
         cache = _cache_evaluator!(model)
@@ -544,6 +551,13 @@ function DiffOpt.forward_differentiate!(model::Model; tol = 1e-6)
     return nothing
 end
 
+MOI.supports(::Model, ::DiffOpt.ReverseDifferentiate) = true
+
+function MOI.set(model::Model, ::DiffOpt.ReverseDifferentiate, value)
+    kws = something(value, (;))
+    return DiffOpt.reverse_differentiate!(model; kws...)
+end
+
 function DiffOpt.reverse_differentiate!(model::Model; tol = 1e-6)
     model.diff_time = @elapsed begin
         cache = _cache_evaluator!(model)

src/QuadraticProgram/QuadraticProgram.jl

@@ -313,7 +313,9 @@ function DiffOpt._get_db(model::Model, ci::EQ)
     return model.back_grad_cache.dν[ci.value]
 end
 
-function DiffOpt.reverse_differentiate!(model::Model)
+MOI.supports(::Model, ::DiffOpt.ReverseDifferentiate) = true
+
+function MOI.set(model::Model, ::DiffOpt.ReverseDifferentiate, ::Nothing)
     model.diff_time = @elapsed begin
         gradient_cache = _gradient_cache(model)
         LHS = gradient_cache.lhs
@@ -354,7 +356,9 @@ end
 struct _QPSets end
 MOI.Utilities.rows(::_QPSets, ci::MOI.ConstraintIndex) = ci.value
 
-function DiffOpt.forward_differentiate!(model::Model)
+MOI.supports(::Model, ::DiffOpt.ForwardDifferentiate) = true
+
+function MOI.set(model::Model, ::DiffOpt.ForwardDifferentiate, ::Nothing)
     model.diff_time = @elapsed begin
         gradient_cache = _gradient_cache(model)
         LHS = gradient_cache.lhs

src/diff_opt.jl

@@ -54,10 +54,12 @@ with respect to the solution set with the [`ReverseVariablePrimal`](@ref) attrib
 The output problem data differentials can be queried with the
 attributes [`ReverseObjectiveFunction`](@ref) and [`ReverseConstraintFunction`](@ref).
 """
-function reverse_differentiate! end
+function reverse_differentiate!(model)
+    return MOI.set(model, ReverseDifferentiate(), nothing)
+end
 
 """
-    forward_differentiate!(model::Optimizer)
+    forward_differentiate!(model::Union{MOI.ModelLike,JuMP.AbstractModel})
 
 Wrapper method for the forward pass.
 This method will consider as input a currently solved problem and
@@ -66,7 +68,9 @@ the [`ForwardObjectiveFunction`](@ref) and  [`ForwardConstraintFunction`](@ref)
 The output solution differentials can be queried with the attribute
 [`ForwardVariablePrimal`](@ref).
 """
-function forward_differentiate! end
+function forward_differentiate!(model)
+    return MOI.set(model, ForwardDifferentiate(), nothing)
+end
 
 """
     empty_input_sensitivities!(model::MOI.ModelLike)
@@ -314,6 +318,40 @@ the differentiation information.
 """
 struct DifferentiateTimeSec <: MOI.AbstractModelAttribute end
 
+"""
+    ReverseDifferentiate <: MOI.AbstractOptimizerAttribute
+
+An `MOI.AbstractOptimizerAttribute` that triggers reverse differentiation
+on the solver. If `MOI.supports(optimizer, DiffOpt.ReverseDifferentiate())`
+returns `true`, then the solver natively supports reverse differentiation
+through the DiffOpt attribute interface, and DiffOpt will delegate
+differentiation directly to the solver instead of using its own
+differentiation backend.
+
+Trigger the computation with:
+```julia
+MOI.set(optimizer, DiffOpt.ReverseDifferentiate(), nothing)
+```
+"""
+struct ReverseDifferentiate <: MOI.AbstractOptimizerAttribute end
+
+"""
+    ForwardDifferentiate <: MOI.AbstractOptimizerAttribute
+
+An `MOI.AbstractOptimizerAttribute` that triggers forward differentiation
+on the solver. If `MOI.supports(optimizer, DiffOpt.ForwardDifferentiate())`
+returns `true`, then the solver natively supports forward differentiation
+through the DiffOpt attribute interface, and DiffOpt will delegate
+differentiation directly to the solver instead of using its own
+differentiation backend.
+
+Trigger the computation with:
+```julia
+MOI.set(optimizer, DiffOpt.ForwardDifferentiate(), nothing)
+```
+"""
+struct ForwardDifferentiate <: MOI.AbstractOptimizerAttribute end
+
 MOI.attribute_value_type(::DifferentiateTimeSec) = Float64
 
 MOI.is_set_by_optimize(::DifferentiateTimeSec) = true
@@ -434,6 +472,14 @@ function MOI.set(model::AbstractModel, ::ForwardObjectiveFunction, objective)
     return
 end
 
+function MOI.supports(::AbstractModel, ::NonLinearKKTJacobianFactorization)
+    return true
+end
+
+function MOI.supports(::AbstractModel, ::AllowObjectiveAndSolutionInput)
+    return true
+end
+
 function MOI.set(
     model::AbstractModel,
     ::NonLinearKKTJacobianFactorization,

src/jump_moi_overloads.jl

@@ -251,6 +251,20 @@ end
 
 MOI.constant(func::IndexMappedFunction) = MOI.constant(func.func)
 
+# Support JuMP.coefficient on plain MOI functions returned by native solvers
+function JuMP.coefficient(
+    func::MOI.ScalarAffineFunction{T},
+    vi::MOI.VariableIndex,
+) where {T}
+    coef = zero(T)
+    for term in func.terms
+        if term.variable == vi
+            coef += term.coefficient
+        end
+    end
+    return coef
+end
+
 function JuMP.coefficient(func::IndexMappedFunction, vi::MOI.VariableIndex)
     return JuMP.coefficient(func.func, func.index_map[vi])
 end

src/moi_wrapper.jl

@@ -548,7 +548,9 @@ end
 
 MOI.get(model::Optimizer, ::ModelConstructor) = model.model_constructor
 
-function reverse_differentiate!(model::Optimizer)
+MOI.supports(::Optimizer, ::ReverseDifferentiate) = true
+
+function MOI.set(model::Optimizer, attr::ReverseDifferentiate, value)
     st = MOI.get(model.optimizer, MOI.TerminationStatus())
     if !in(st, (MOI.LOCALLY_SOLVED, MOI.OPTIMAL))
         error(
@@ -562,17 +564,22 @@ function reverse_differentiate!(model::Optimizer)
                   "Set `DiffOpt.AllowObjectiveAndSolutionInput()` to `true` to silence this warning."
         end
     end
-    diff = _diff(model)
-    MOI.set(
-        diff,
-        NonLinearKKTJacobianFactorization(),
-        model.input_cache.factorization,
-    )
-    MOI.set(
-        diff,
-        AllowObjectiveAndSolutionInput(),
-        model.input_cache.allow_objective_and_solution_input,
-    )
+    diff = _diff(model, attr)
+    # Native differentiation interface are allow to not support it
+    if MOI.supports(diff, NonLinearKKTJacobianFactorization())
+        MOI.set(
+            diff,
+            NonLinearKKTJacobianFactorization(),
+            model.input_cache.factorization,
+        )
+    end
+    if MOI.supports(diff, AllowObjectiveAndSolutionInput())
+        MOI.set(
+            diff,
+            AllowObjectiveAndSolutionInput(),
+            model.input_cache.allow_objective_and_solution_input,
+        )
+    end
     for (vi, value) in model.input_cache.dx
         MOI.set(diff, ReverseVariablePrimal(), model.index_map[vi], value)
     end
@@ -593,7 +600,7 @@ function reverse_differentiate!(model::Optimizer)
             end
         end
     end
-    return reverse_differentiate!(diff)
+    return MOI.set(diff, attr, value)
 end
 
 # Gradient evaluation functions for objective sensitivity fallbacks
@@ -639,13 +646,12 @@ function _eval_gradient(
 end
 
 function _fallback_set_reverse_objective_sensitivity(model::Optimizer, val)
-    diff = _diff(model)
     obj_type = MOI.get(model, MOI.ObjectiveFunctionType())
     obj_func = MOI.get(model, MOI.ObjectiveFunction{obj_type}())
     grad = _eval_gradient(model, obj_func)
     for (xi, df_dxi) in grad
         MOI.set(
-            diff,
+            model.diff,
             ReverseVariablePrimal(),
             model.index_map[xi],
             df_dxi * val,
@@ -666,24 +672,30 @@ function _copy_forward_in_constraint(diff, index_map, con_map, constraints)
     return
 end
 
-function forward_differentiate!(model::Optimizer)
+MOI.supports(model::Optimizer, attr::ForwardDifferentiate) = true
+
+function MOI.set(model::Optimizer, attr::ForwardDifferentiate, value)
     st = MOI.get(model.optimizer, MOI.TerminationStatus())
     if !in(st, (MOI.LOCALLY_SOLVED, MOI.OPTIMAL))
         error(
             "Trying to compute the forward differentiation on a model with termination status $(st)",
         )
     end
-    diff = _diff(model)
-    MOI.set(
-        diff,
-        NonLinearKKTJacobianFactorization(),
-        model.input_cache.factorization,
-    )
-    MOI.set(
-        diff,
-        AllowObjectiveAndSolutionInput(),
-        model.input_cache.allow_objective_and_solution_input,
-    )
+    diff = _diff(model, attr)
+    if MOI.supports(diff, NonLinearKKTJacobianFactorization())
+        MOI.set(
+            diff,
+            NonLinearKKTJacobianFactorization(),
+            model.input_cache.factorization,
+        )
+    end
+    if MOI.supports(diff, AllowObjectiveAndSolutionInput())
+        MOI.set(
+            diff,
+            AllowObjectiveAndSolutionInput(),
+            model.input_cache.allow_objective_and_solution_input,
+        )
+    end
     T = Float64
     list = MOI.get(
         model,
@@ -700,7 +712,7 @@ function forward_differentiate!(model::Optimizer)
                 MOI.Parameter(value),
             )
         end
-        return forward_differentiate!(diff)
+        return MOI.set(diff, attr, value)
     end
     # @show "func mode"
     if model.input_cache.objective !== nothing
@@ -733,7 +745,7 @@ function forward_differentiate!(model::Optimizer)
             model.input_cache.vector_constraints[F, S],
         )
     end
-    return forward_differentiate!(diff)
+    return MOI.set(diff, attr, value)
 end
 
 function empty_input_sensitivities!(model::Optimizer)
@@ -782,8 +794,24 @@ function _instantiate_diff(model::Optimizer, constructor)
     return model_bridged
 end
 
-function _diff(model::Optimizer)
-    if model.diff === nothing
+# Find the native differentiation solver in the optimizer chain.
+# Cached in `model.diff` to avoid repeated unwrapping.
+function _native_diff_solver(model::Optimizer)
+    if isnothing(model.diff)
+        model.diff = model.optimizer
+        model.index_map = MOI.Utilities.identity_index_map(model.optimizer)
+    end
+    return model.diff
+end
+
+function _diff(
+    model::Optimizer,
+    attr::Union{ForwardDifferentiate,ReverseDifferentiate},
+)
+    if MOI.supports(model.optimizer, attr)
+        model.diff = model.optimizer
+        model.index_map = MOI.Utilities.identity_index_map(model.optimizer)
+    elseif isnothing(model.diff)
         _check_termination_status(model)
         model_constructor = MOI.get(model, ModelConstructor())
         if isnothing(model_constructor)
@@ -1128,15 +1156,7 @@ function MOI.get(model::Optimizer, attr::DifferentiateTimeSec)
     return MOI.get(model.diff, attr)
 end
 
-function MOI.supports(
-    ::Optimizer,
-    ::NonLinearKKTJacobianFactorization,
-    ::Function,
-)
-    return true
-end
-
-function MOI.supports(::Optimizer, ::AllowObjectiveAndSolutionInput, ::Bool)
+function MOI.supports(::Optimizer, ::NonLinearKKTJacobianFactorization)
     return true
 end
 

src/parameters.jl

@@ -479,8 +479,11 @@ function empty_input_sensitivities!(model::POI.Optimizer{T}) where {T}
     return
 end
 
-function forward_differentiate!(model::POI.Optimizer{T}) where {T}
-    empty_input_sensitivities!(model.optimizer)
+function MOI.set(
+    model::POI.Optimizer{T},
+    attr::ForwardDifferentiate,
+    value,
+) where {T}
     ctr_types = MOI.get(model, POI.ListOfParametricConstraintTypesPresent())
     for (F, S, P) in ctr_types
         dict = MOI.get(
@@ -497,7 +500,7 @@ function forward_differentiate!(model::POI.Optimizer{T}) where {T}
     elseif obj_type <: POI.ParametricCubicFunction
         _cubic_objective_set_forward!(model)
     end
-    forward_differentiate!(model.optimizer)
+    MOI.set(model.optimizer, attr, value)
     return
 end
 
@@ -699,8 +702,8 @@ function _quadratic_objective_get_reverse!(model::POI.Optimizer{T}) where {T}
     return
 end
 
-function reverse_differentiate!(model::POI.Optimizer)
-    reverse_differentiate!(model.optimizer)
+function MOI.set(model::POI.Optimizer, attr::ReverseDifferentiate, value)
+    MOI.set(model.optimizer, attr, value)
     sensitivity_data = _get_sensitivity_data(model)
     empty!(sensitivity_data.parameter_output_backward)
     sizehint!(

test/jump.jl

@@ -688,11 +688,9 @@ function test_conic_supports()
         DiffOpt.ForwardConstraintSet(),
         MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{Float64}},
     )
-    function some end
     @test MOI.supports(
         backend(model),
         DiffOpt.NonLinearKKTJacobianFactorization(),
-        some,
     )
     MOI.is_set_by_optimize(DiffOpt.ReverseConstraintFunction())
     MOI.is_set_by_optimize(DiffOpt.ReverseConstraintSet())