MadDiff

MadDiff implements forward and reverse mode implicit differentiation for MadSuite solvers. MadDiff leverages MadNLP's modular KKT and linear solver infrastructure, supporting LP, QP, and NLP using KKT systems from MadNLP, MadIPM, MadNCL, and HybridKKT.

Warning

MadDiff is a work-in-progress and requires installing forks of several dependencies. Proceed with caution and verify correctness before use.

NLPModels interface

Note

The NLPModels interface requires that your AbstractNLPModel implementation includes the parametric AD API. Currently, this is automated only for the case when using MadNLP through JuMP (fork) or when using ExaModels (fork); support for other solvers and modelers is planned.

nlp = ...  # must implement parametric AD API
solver = MadNLP.MadNLPSolver(nlp)
solution = MadNLP.solve!(solver)

diff = MadDiff.MadDiffSolver(solver)

dL_dx, dL_dy, dL_dzl, dL_dzu = ...  # loss sensitivity vectors
rev = MadDiff.vector_jacobian_product!(diff; dL_dx, dL_dy, dL_dzl, dL_dzu)
rev.grad_p  # gradient of the loss with respect to the parameters

JuMP interface

MadDiff aims to be a drop-in replacement for DiffOpt with MadNLP. Simply switch DiffOpt.diff_model(MadNLP.Optimizer) for MadDiff.diff_model(MadNLP.Optimizer) and enjoy the speedup!

using JuMP, DiffOpt
using MadDiff, MadNLP

model = MadDiff.diff_model(MadNLP.Optimizer)
@variable(model, x)
@variable(model, p in MOI.Parameter(1.0))
@constraint(model, x >= 2p)
@objective(model, Min, x^2)
optimize!(model)

DiffOpt.empty_input_sensitivities!(model)
MOI.set(model, DiffOpt.ForwardConstraintSet(), ParameterRef(p), MOI.Parameter(1.0))
DiffOpt.forward_differentiate!(model)
dx = MOI.get(model, DiffOpt.ForwardVariablePrimal(), x)

DiffOpt.empty_input_sensitivities!(model)
MOI.set(model, DiffOpt.ReverseVariablePrimal(), x, 1.0)
DiffOpt.reverse_differentiate!(model)
dp = MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(p)).value