Adds support for differentable-native solvers

[PTNobel]

Hey JuMP team!

I'm working on moreau.so @optimalintellect; and we natively support differentiable optimization, and can reuse some data structures between the forward and backward pass. However, to do that, we need to preserve a handle that points to the data structures inbetween the forward and backward pass.

I wanted to support DiffOpt natively in our Julia frontend, so I am adding support in DiffOpt for an MOI interface to declare that it has native differentiation support and for that to be used.

I also added documentation for this.

This is my first time using Julia, so Claude Code was invaluable. I have read all of the code, and I believe I understand it. That said, I make no promises that I have any idea what the Julia ecosystem expects, etc. so please suggest as many or as extensive of revisions as desired.

[klamike]

FYI, I have something similar for MadDiff's DiffOpt integration, see https://github.com/MadNLP/MadDiff.jl/blob/main/ext/DiffOptExt/DiffOptExt.jl and https://github.com/MadNLP/MadDiff.jl/tree/main/ext/MathOptInterfaceExt

The biggest difference is that there, I implement a ModelConstructor instead, since I wanted to not touch the internals of DiffOpt. It requires users to either set the attribute themselves, or swap DiffOpt.diff_model for MadDiff.diff_model which takes care of that. There are some other differences to do with the implementation of the components, but the high-level idea is the same as this PR.

Curious what the maintainers think. Not sure about how to best implement it, but it would be great to have this capability officially supported in DiffOpt!

[joaquimg]

Hey,

Thanks for the interest in plugging this into DiffOpt!

I will take a close look over the weekend.

I must confess we did not give much thought to the native diff solver, but this is past due.

I'd say that simply reusing the DiffOpt API should be relatively easy: simply overload some getters and setters in your MOI interface.
The issue with that approach is not taking advantage of:
1 - the MOI bridging system
2 - POI parameters handling (If your solver natively supports parameters, then this is not an issue).

I have this (very) experimental work here: https://github.com/jump-dev/DiffOpt.jl/tree/jg/lpbasis
For LP solvers diff. This approach attempts to keep (1) and (2) working.

Starting point:
Before working with DiffOpt, however, we should make sure moreau.so works smoothly with MOI.
I basically mean a MOI interface (tipically in a Moreau.jl julia package): https://jump.dev/MathOptInterface.jl/stable/tutorials/implementing/ (beware of the danger note!)

If that already exists, send me a pointer!

[PTNobel]

The Moreau.jl package is working and passing the MOI.Test. It's waiting for approval to be made public, but right now it's very similar to SCS.jl internally

[joaquimg]

Will Moreau.jl natively support parameters?

[PTNobel]

I am not sure; what MOI features would that require?

[joaquimg]

My suggestion is:
Only try supporting parameters natively if the inner solver does that.
Do not add the complexity of handling parameter to the MOI wrapper level, we can use a ParametricOptInterface layer to enable parameters.

If the inner solver supports parameters the way to go is supporting the MOI.Parameter set, when implementing MOI.add_constrained_variable so that the variable is a parameter for the solver perspective.
Ipopt.jl is possibly the only example with such native support.
We can also help with details when Moreau.jl is out.

[odow]

I discussed this with Steven.

1. First, we should make Moreau.jl public.
2. Then, we should add whatever features needed to Moreau.jl to make Moreau's differentiable API accessible
3. Then, we should compare what is in DiffOpt and what is in Moreau and move any common attributes to MOI

I don't think Moreau needs to know about DiffOpt. DiffOpt is effectively a layer for solvers that don't have native differentiability.

[joaquimg]

I do not agree with that order of priorities.

Also, that is not entirely true. DiffOpt is both interface defining and a helper for solvers that do not support differentiability.

We did discuss differentiability related attributes in MOI before and the answer was not to touch MOI. If you changed your mind, I think we are ready to push differentiation attributes to MOI. I'd be happy with this.

However, I fear that diff attributes in MOI might be a long discussion since there are many important details.

About the needed attributes:
We need the attributes defined here in DiffOpt:
ForwardConstraintFunction (for variable coefficients and for constants in constraints and objective)
ForwardConstraintSet (for parameters)
ForwardVariablePrimal
ReverseVariablePrimal
ReverseConstraintSet (for parameters)
ReverseConstraintFunction

This two are extremely useful in practice:
ForwardObjectiveSensitivity
ReverseObjectiveSensitivity

moreover, we need:
reverse_differentiate!
forward_differentiate!
empty_input_sensitivities!

In DiffOpt we have been making sure everything is always supported, but we can certainly have some supports... type functions.

Moreau supports differentiability more explicitly.

Houwever, we already have solver with partial support for that that we are not currently exploiting: HiGHS, Gurobi, Xpress (possibly CPLEX).
All of these could easily support forward and backward modes with respect to RHS. They simply give a different name to it. A proof of concept is in a branch referenced in the comments above.

Extremely important things to consider:
1 - Bridges are important
2 - Parameters are very important (it is excruciatingly painful to work with this without parameters)

Therefore, we have absolutely no need to wait for Moreau to discuss the MOI differentiation API.

I see absolutely no reason to push a full Moreau MOI-level implementation of differentiation without looking at DiffOpt. Sure, we can see the low-level API, but there is no reason to duplicate work and do it twice (a "first try" and a MOI review)

I agree that having an API defined at the MOI level would be very helpful, as we could potentially make Bridges and POI explicitly aware of it.

But again, I do not think of this as a quick endeavor.

Waiting for all that will simply not allow Moreau to be used cleanly in JuMP, whereas DiffOpt would allow it quickly.

So my counter proposal is to have parallel non-blocking efforts:

• Make Moreau public
• Start a serious discussion about moving such attributes do MOI (and JuMP)
• Review changes in DiffOpt to support external solvers (if the bullet above moves fast, which I doubt, I will happily deprecate this)

[blegat]

We shouldn't need that, the goal of attributes is to flow through the layers so that you never need to unwrap.

[blegat]

I agree with Joaquim, I don't see why we should move things to MOI. MOI has reached v1 so anything we move there is frozen. Adding DiffOpt as dependency to Moreau or even as a packages extension is preferrable.
We can just start dicussing this PR directly. DiffOpt was designed in a way to make it possible for solvers to give their derivatives directly, at the time we were thinking of doing this for OSQP and waiting for it to be part of the released version of osqp but never got to do it so I'm curious to see how to make it work with Moreau. Of course, it will be made easier once it's public but no need to wait.

I'm wondering if it wouldn't be possible to do something much simpler that what Claude Code suggested here.
I like the idea of defining

struct BackwardDifferentiate <: MOI.AbstractModelAttribute end

and then checking whether the solver supports it with

MOI.supports(model.optimizer, DiffOpt.BackwardDifferentiate())

and if it does just triggering the computation with

MOI.set(model.optimizer, DiffOpt.BackwardDifferentiate())

However, for setting the attributes, it shouldn't do the _unwrap etc... It should just pass the attributes we defined for it.

So in

DiffOpt.jl/src/moi_wrapper.jl

Line 551 in c010b53

function reverse_differentiate!(model::Optimizer)

in case MOI.supports(model.optimizer, DiffOpt.BackwardDifferentiate()) returns true, then we copy the input_cache directly into model.optimizer, not model.diff. So maybe what we can simply do is make

DiffOpt.jl/src/moi_wrapper.jl

Line 785 in c010b53

function _diff(model::Optimizer)

return model.optimizer and model.diff can just be nothing, it won't be used.
These parameters are all passed down through the layers, and transformed accordingly (as joaquim says we implement transformation in the bridges on a per-need basis, it's easy but we only do it for a bridge when we need it, don't hesitate to let us know if it's needed for a bridge your solver is using).
Then, we request the differentiation to the solver with

MOI.set(model.optimizer, DiffOpt.BackwardDifferentiate())

Here, to make sure it's passed down to the solver and not the cache, maybe it should be an AbstractOptimizerAttribute, not an AbstractModelAttribute.
The, we should make

DiffOpt.jl/src/moi_wrapper.jl

Line 839 in c010b53

function _checked_diff(model::Optimizer, attr::MOI.AnyAttribute, call)

return model.optimizer, and all results will be queried from the solver as the attributes defined in DiffOpt, and transformed across the layers as expected.
This was designed with attributes for this use case to work, I just never tried it with an actual solver so there were these remaining pieces to make it work but I'd prefer to first try it this way, which should be much simpler in my opinion. If there are good reasons that this approach doesn't work then we can consider the approach of this PR (I wouldn't want Claude Code to feel offended :) )