From fd7978a5b5c0fd849900353f244bda86f1341b3f Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Thu, 19 Feb 2026 22:26:42 -0500 Subject: [PATCH] initial vno support --- .../MathOptInterfaceExt.jl | 4 +++ ext/MathOptInterfaceExt/forward_mode.jl | 32 +++++++++++++++++-- ext/MathOptInterfaceExt/moi_wrapper.jl | 3 ++ ext/MathOptInterfaceExt/reverse_mode.jl | 22 +++++++++++++ test/Project.toml | 2 +- 5 files changed, 60 insertions(+), 3 deletions(-) diff --git a/ext/MathOptInterfaceExt/MathOptInterfaceExt.jl b/ext/MathOptInterfaceExt/MathOptInterfaceExt.jl index 16ec8ca..7c7105b 100644 --- a/ext/MathOptInterfaceExt/MathOptInterfaceExt.jl +++ b/ext/MathOptInterfaceExt/MathOptInterfaceExt.jl @@ -9,24 +9,28 @@ mutable struct ForwardModeData{T} param_perturbations::Dict{MOI.ConstraintIndex, T} primal_sensitivities::Dict{MOI.VariableIndex, T} dual_sensitivities::Dict{MOI.ConstraintIndex, T} + vector_dual_sensitivities::Dict{MOI.ConstraintIndex, Vector{T}} objective_sensitivity::T end ForwardModeData{T}() where {T} = ForwardModeData{T}( Dict{MOI.ConstraintIndex, T}(), Dict{MOI.VariableIndex, T}(), Dict{MOI.ConstraintIndex, T}(), + Dict{MOI.ConstraintIndex, Vector{T}}(), zero(T), ) mutable struct ReverseModeData{T} primal_seeds::Dict{MOI.VariableIndex, T} dual_seeds::Dict{MOI.ConstraintIndex, T} + vector_dual_seeds::Dict{MOI.ConstraintIndex, Vector{T}} param_outputs::Dict{MOI.ConstraintIndex, T} dobj::Union{Nothing, T} end ReverseModeData{T}() where {T} = ReverseModeData{T}( Dict{MOI.VariableIndex, T}(), Dict{MOI.ConstraintIndex, T}(), + Dict{MOI.ConstraintIndex, Vector{T}}(), Dict{MOI.ConstraintIndex, T}(), nothing, ) diff --git a/ext/MathOptInterfaceExt/forward_mode.jl b/ext/MathOptInterfaceExt/forward_mode.jl index b276c13..174d66a 100644 --- a/ext/MathOptInterfaceExt/forward_mode.jl +++ b/ext/MathOptInterfaceExt/forward_mode.jl @@ -54,7 +54,7 @@ function _forward_differentiate_impl!(model::Optimizer{OT, T}) where {OT, T} dy = _get_dy_cache!(model, n_con) dy .= (.-obj_sign) .* dy_cpu - _store_dual_sensitivities!(model.forward.dual_sensitivities, inner, dy) + _store_dual_sensitivities!(model.forward.dual_sensitivities, model.forward.vector_dual_sensitivities, inner, dy) _store_bound_dual_sensitivities!(model, sens, result, inner) model.forward.objective_sensitivity = result.dobj[] return @@ -68,10 +68,30 @@ function _constraint_row(inner, ci::MOI.ConstraintIndex{F, S}) where {F, S} end end -function _store_dual_sensitivities!(dual_sensitivities, inner, dy) +function _vno_rows( + inner, + ci::MOI.ConstraintIndex{MOI.VectorOfVariables, MOI.VectorNonlinearOracle{Float64}}, +) + offset = length(inner.qp_data) + for i in 1:(ci.value - 1) + _, s = inner.vector_nonlinear_oracle_constraints[i] + offset += s.set.output_dimension + end + _, s = inner.vector_nonlinear_oracle_constraints[ci.value] + return offset .+ (1:s.set.output_dimension) +end + +function _store_dual_sensitivities!(dual_sensitivities, vector_dual_sensitivities, inner, dy) for (F, S) in MOI.get(inner, MOI.ListOfConstraintTypesPresent()) F == MOI.VariableIndex && continue S <: MOI.Parameter && continue + if F == MOI.VectorOfVariables && S == MOI.VectorNonlinearOracle{Float64} + for ci in MOI.get(inner, MOI.ListOfConstraintIndices{F, S}()) + rows = _vno_rows(inner, ci) + vector_dual_sensitivities[ci] = dy[rows] + end + continue + end for ci in MOI.get(inner, MOI.ListOfConstraintIndices{F, S}()) row = _constraint_row(inner, ci) dual_sensitivities[ci] = dy[row] @@ -127,6 +147,14 @@ function MOI.get(model::Optimizer, ::MadDiff.ForwardConstraintDual, ci::MOI.Cons return model.forward.dual_sensitivities[ci] end +function MOI.get( + model::Optimizer, + ::MadDiff.ForwardConstraintDual, + ci::MOI.ConstraintIndex{MOI.VectorOfVariables, MOI.VectorNonlinearOracle{Float64}}, +) + return model.forward.vector_dual_sensitivities[ci] +end + function MOI.get(model::Optimizer, ::MadDiff.ForwardObjectiveSensitivity) return model.forward.objective_sensitivity end diff --git a/ext/MathOptInterfaceExt/moi_wrapper.jl b/ext/MathOptInterfaceExt/moi_wrapper.jl index d692e1d..f70abe2 100644 --- a/ext/MathOptInterfaceExt/moi_wrapper.jl +++ b/ext/MathOptInterfaceExt/moi_wrapper.jl @@ -81,6 +81,7 @@ function MOI.empty!(m::Optimizer) empty!(m.forward.param_perturbations) empty!(m.reverse.primal_seeds) empty!(m.reverse.dual_seeds) + empty!(m.reverse.vector_dual_seeds) m.reverse.dobj = nothing return _invalidate_sensitivity!(m) end @@ -103,6 +104,7 @@ function MadDiff.empty_input_sensitivities!(model::Optimizer) empty!(model.forward.param_perturbations) empty!(model.reverse.primal_seeds) empty!(model.reverse.dual_seeds) + empty!(model.reverse.vector_dual_seeds) model.reverse.dobj = nothing return _clear_outputs!(model) end @@ -110,6 +112,7 @@ end function _clear_outputs!(m::Optimizer{OT, T}) where {OT, T} empty!(m.forward.primal_sensitivities) empty!(m.forward.dual_sensitivities) + empty!(m.forward.vector_dual_sensitivities) m.forward.objective_sensitivity = zero(T) empty!(m.reverse.param_outputs) return m.diff_time = zero(T) diff --git a/ext/MathOptInterfaceExt/reverse_mode.jl b/ext/MathOptInterfaceExt/reverse_mode.jl index 42360c4..8baca87 100644 --- a/ext/MathOptInterfaceExt/reverse_mode.jl +++ b/ext/MathOptInterfaceExt/reverse_mode.jl @@ -18,6 +18,16 @@ function MOI.set( return _clear_outputs!(model) # keep KKT factorization end +function MOI.set( + model::Optimizer, + ::MadDiff.ReverseConstraintDual, + ci::MOI.ConstraintIndex{MOI.VectorOfVariables, MOI.VectorNonlinearOracle{Float64}}, + value::AbstractVector, + ) + model.reverse.vector_dual_seeds[ci] = value + return _clear_outputs!(model) # keep KKT factorization +end + function MOI.set( model::Optimizer{OT, T}, ::MadDiff.ReverseObjectiveSensitivity, @@ -73,6 +83,15 @@ function _process_reverse_dual_input!( dL_dy[row] = val end +function _process_reverse_dual_input!( + ci::MOI.ConstraintIndex{MOI.VectorOfVariables, MOI.VectorNonlinearOracle{Float64}}, + val::AbstractVector, + inner, dL_dy, dL_dzl, dL_dzu, +) + rows = _vno_rows(inner, ci) + dL_dy[rows] .= val +end + function _reverse_differentiate_impl!(model::Optimizer{OT, T}) where {OT, T} inner = model.inner solver = inner.solver @@ -99,6 +118,9 @@ function _reverse_differentiate_impl!(model::Optimizer{OT, T}) where {OT, T} for (ci, val) in model.reverse.dual_seeds _process_reverse_dual_input!(ci, val, inner, dL_dy, dL_dzl, dL_dzu) end + for (ci, val) in model.reverse.vector_dual_seeds + _process_reverse_dual_input!(ci, val, inner, dL_dy, dL_dzl, dL_dzu) + end dL_dy .*= -solver.cb.obj_sign dobj = model.reverse.dobj diff --git a/test/Project.toml b/test/Project.toml index 92072aa..6319113 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -21,7 +21,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] -DiffOpt = {rev = "mk/337and338", url = "https://github.com/klamike/DiffOpt.jl.git"} +DiffOpt = {path = "../../DiffOpt.jl"} HybridKKT = {rev = "mk/latest", url = "https://github.com/klamike/HybridKKT.jl.git"} MadDiff = {path = ".."} MadIPM = {rev = "mk/flip", url = "https://github.com/klamike/MadIPM.jl.git"}