Where is JuMP's sparse AD code, and can DifferentiationInterface help?

[gdalle]

Hi there @odow and friends!

Looking at our work on sparse autodiff with @adrhill and @amontoison, I've been wondering how it could be useful to JuMP. We have developed a combination of three new packages:

• for sparsity detection (SparseConnectivityTracer.jl)
• for matrix colorings and decompression (SparseMatrixColorings.jl)
• for backend-independent gradients, (sparse) Jacobians and (sparse) Hessians (DifferentiationInterface.jl)

From what I understand, JuMP's current sparse AD engine is Nonlinear.ReverseAD, but there are also experiments going on in MathOptSymbolicAD.jl. Did I miss anything?

I'm not suggesting that DifferentiationInterface and friends should replace your own default AD solution. But since you say in the docs that testing other AD backends would be nice, maybe there's an angle there?
In any case, perhaps DI could be a nice addition to the docs page on autodiff of user-defined operators?

[odow]

Did I miss anything?

Nope. That is the situation. All nonlinear/AD stuff is in the MOI.Nonllinear module. There is nothing in JuMP.

But since you say in the docs that testing other AD backends would be nice, maybe there's an angle there?

New AD backends need to be implemented as an oracle like this:

https://github.com/lanl-ansi/MathOptSymbolicAD.jl/blob/e5d3eaa1f56346d53a5ef8f895743ab3e89ff38c/src/nonlinear_oracle.jl#L278-L496

In any case, perhaps DI could be a nice addition to the docs page on autodiff of user-defined operators?

Feel free to open a PR 😄

[amontoison]

What could be useful for JuMP is star coloring for the Hessian of the Lagrangian.
It requires more colors compared to acyclic coloring (the coloring used by default in MOI.Nonlinear), but it puts less pressure on memory by allowing not storing all columns of the compressed Hessian for very large problems.

The issue is that the code related to nonlinear stuff in MOI is highly specialized for what they use in JuMP, so it's probably not easy to create a bridge.

[odow]

I have no plans to work on a DI interface, but anyone interested is encouraged to give it a go. It can (and should) be done first in a separate repo like MathOptSymbolicAD.jl. One goal of refactoring the nonlinear stuff out of JuMP was so that we could experiment without needing to make source code changes to JuMP or MathOptInterface.

For the user-defined gradients with DI, I'm open to including this in the documentation, but it is lower priority, so perhaps we could close this issue in favor of adding a TODO item to #2348.

[amontoison]

I agree with that Oscar.

[odow]

Closing in favor of #2348 (comment)