Skip to content

Interface for solver supporting diff natively #356

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
blegat merged 24 commits into from
Apr 7, 2026
Merged

Interface for solver supporting diff natively #356

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
5ce30a7
Interface for solver supporting diff natively
blegat Mar 25, 2026
5f8ac99
Fixes
blegat Mar 26, 2026
8eb0c2b
Fix format
blegat Mar 26, 2026
f29056b
Claude improves coverage
blegat Mar 26, 2026
cabc199
Fix format
blegat Mar 26, 2026
463bcda
Add tests
blegat Mar 26, 2026
c49c519
fix format
blegat Mar 26, 2026
e3c8210
Rename
blegat Apr 1, 2026
c8e72f4
Use the attributes in inner models too
blegat Apr 1, 2026
03f4af7
Simplify
blegat Apr 1, 2026
b4aee93
Fix format
blegat Apr 1, 2026
6e25d00
Fix
blegat Apr 1, 2026
b75c198
Restore keyword option for nonlinear
blegat Apr 1, 2026
cbff8af
Fix
blegat Apr 1, 2026
7f652bc
Fix
blegat Apr 1, 2026
af656ae
Fixes
blegat Apr 1, 2026
fe90ee8
Fix
blegat Apr 1, 2026
3e49830
Fix
blegat Apr 3, 2026
1332d5a
Fix
blegat Apr 3, 2026
c8fe60c
Fix
blegat Apr 3, 2026
890fd48
Simpler
blegat Apr 3, 2026
2316a7b
Fix format
blegat Apr 3, 2026
f6f3357
Simplify handling of POI layer
blegat Apr 3, 2026
49f824e
Fix format
blegat Apr 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions src/ConicProgram/ConicProgram.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
14 changes: 14 additions & 0 deletions src/NonLinearProgram/NonLinearProgram.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
8 changes: 6 additions & 2 deletions src/QuadraticProgram/QuadraticProgram.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
52 changes: 49 additions & 3 deletions src/diff_opt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
14 changes: 14 additions & 0 deletions src/jump_moi_overloads.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Comment on lines +255 to +258
Copy link
Copy Markdown
Member

@joaquimg joaquimg Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

more type piracy...

Copy link
Copy Markdown
Member Author

@blegat blegat Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's maybe not needed in the end, let me check

Copy link
Copy Markdown
Member Author

@blegat blegat Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See #360

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
Expand Down
100 changes: 60 additions & 40 deletions src/moi_wrapper.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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())
Comment on lines +675 to 678
Copy link
Copy Markdown
Member

@joaquimg joaquimg Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there reasonable precedence for this?

Copy link
Copy Markdown
Member Author

@blegat blegat Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using MOI.set to trigger a computation ? I don't think so but I think it's the most layer-friendly approach

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,
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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

Expand Down
13 changes: 8 additions & 5 deletions src/parameters.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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

Expand Down Expand Up @@ -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!(
Expand Down
2 changes: 0 additions & 2 deletions test/jump.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
Loading
Loading