Nonlinear reverse: single adjoint solve instead of dense full-parameter sensitivity

[yeabbratz]

Motivation

For NonLinearProgram models, reverse_differentiate! currently forms the dense parameter
sensitivity ∂s = -K⁻¹N (kkt_dim × P — one forward solve per parameter) and only then contracts
it with the seed. A reverse request is a vector–Jacobian product, so only one adjoint vector is
needed: the dense solve is O(P) wasted time and memory. On a downstream differentiable-ACOPF
project this dominated the backward pass (12 s / 20 GiB per backward at 3633 parameters) and
blocked large instances.

Change

reverse_differentiate!(::NonLinearProgram.Model) now computes

z  = Kᵀ \ Δw          # ONE right-hand side, same factorization (and same inertia correction)
Δp = -Nᵀ z            # sparse matvec

via a new _compute_sensitivity_adjoint in nlp_utilities.jl, instead of materializing ∂s.
Details that keep it exactly the quantity the dense path contracted:

• the JuMP-convention sign adjustments that _compute_sensitivity applies to the dual rows
  of ∂s are applied to the corresponding entries of the seed (algebraically identical);
• the objective seed folds into the same right-hand side (df_dp = ∇ₓfᵀ ∂s[1:n,:] + ∇ₚfᵀ, so
  its ∂s-dependent part rides along as ∇ₓf · dobj added to the primal seed entries, and the
  direct ∇ₚf · dobj term is added after the solve) — mixed dx/dy/dobj requests still take
  one adjoint solve;
• the singular/inertia-correction fallback is unchanged (K === nothing ⇒ zero solve component);
• forward_differentiate! and _compute_sensitivity are unchanged; no public API change, no
  new flag — the reverse path is a VJP by definition, so it always benefits.

Results

In-tree reproducer (script below; sparse parameterized NLP, chain-coupled quartic objective + one
nonlinear inequality per variable; Ipopt; median of 5, M4 Pro):

P	KKT dim	adjoint (this PR)	dense (pre-PR algorithm)	ratio	max rel Δp diff
50	203	0.6 ms / 0.001 GiB	0.4 ms / 0.001 GiB	~1×	4.9e-16
200	803	1.5 ms / 0.006 GiB	1.9 ms / 0.007 GiB	1.3×	4.8e-16
800	3203	5.5 ms / 0.027 GiB	26.4 ms / 0.097 GiB	4.8×	1.5e-15
2000	8003	15.8 ms / 0.090 GiB	164.8 ms / 0.577 GiB	10.5×	1.2e-15

(The "dense" column runs the pre-PR contraction via the still-shipped _compute_sensitivity, so
the comparison is exact; at small P the two are within noise.)

Downstream benchmark that motivated this (parameterized ACOPF, PGLib cases, Ipopt, DiffOpt 0.6.0
with this change applied as a local override): per backward at case2000_goc__api (P = 3633)
9.6 s / 20.2 GiB → 0.94 s / 0.70 GiB (×10.2 time, ×29 allocation); the speedup grows with P
(×1.1 / ×1.2 / ×2.3 / ×3.6 / ×10.2 across case14/30/179/500/2000). Reverse allocation now tracks
KKT assembly, not KKT×P.

Correctness

Same quantity, Δp = -Nᵀ K⁻ᵀ Δw; the only difference from the dense path is floating-point solve
order (one transposed solve instead of P forward columns), so results are numerically
equivalent but not bitwise — ≤1.5e-15 relative on the reproducer above, ~1e-11 relative on the
ill-conditioned ACOPF KKT systems downstream (≈4 orders below Ipopt's 1e-8 tolerance, and
validated there against central finite differences at two operating points). In-tree:

• the full existing test suite passes unchanged — including the NLP reverse tests with dx, dy
  and dobj seeds and the FiniteDiff-validated ones, which now exercise the adjoint path;
• a new test (test_reverse_adjoint_matches_forward_jvp) cross-checks reverse (VJP) against
  forward (JVP) component-wise, ⟨Δw, ∂x/∂pᵢ⟩ + dobj·∂f/∂pᵢ == Δpᵢ, across parameter counts and
  with simultaneous primal + objective seeds.

Reproducer

bench_reverse_adjoint.jl (standalone; needs JuMP + Ipopt + this branch)

# bench_reverse_adjoint.jl — standalone reproducer for the DiffOpt PR "Nonlinear reverse:
# single adjoint solve instead of dense full-parameter sensitivity".
#
# Self-contained (no project-specific dependencies): needs DiffOpt (the PR branch), JuMP, Ipopt.
#   julia -e 'import Pkg; Pkg.activate(temp=true); Pkg.add(["JuMP","Ipopt"]);
#             Pkg.develop(path="path/to/DiffOpt.jl"); include("bench_reverse_adjoint.jl")'
#
# For each parameter count P it builds a sparse parameterized NLP (chain-coupled quartic
# objective, one nonlinear inequality per variable, one aggregate inequality), solves with Ipopt,
# then times per backward:
#   * reverse — `DiffOpt.reverse_differentiate!` (on the PR branch: ONE adjoint solve
#               `Δp = -Nᵀ(K⁻ᵀ Δw)`);
#   * dense   — the pre-PR algorithm, reproduced from the still-shipped forward machinery:
#               `∂s, _ = _compute_sensitivity(model)` then `Δp = ∂sᵀ Δw` (this IS master's
#               reverse contraction, so the comparison is exact on either checkout). Note the
#               dense timing does NOT include the evaluator rebuild that `reverse` pays via
#               `_cache_evaluator!` — the asymmetry is conservative (it understates the gap).
# and reports the gradient match. Expected: identical Δp to ~1e-11 relative (floating-point
# solve-order change only: one transposed RHS instead of P forward columns).

using JuMP
using Ipopt
using DiffOpt
import MathOptInterface as MOI
using Printf
using Statistics

const NLP = DiffOpt.NonLinearProgram
const REPS = 5

function build_and_solve(P)
    model = DiffOpt.nonlinear_diff_model(Ipopt.Optimizer)
    set_silent(model)
    @variable(model, p[i = 1:P] in MOI.Parameter(1.0 + 0.1 * sin(i)))
    @variable(model, x[1:P])
    @constraint(model, [i = 1:P], x[i] * (1 + 0.1 * sin(p[i])) >= p[i])
    @constraint(model, sum(x) >= 0.6 * P)
    @objective(model, Min,
        sum((x[i] - p[i])^2 for i in 1:P) +
        0.01 * sum(x[i]^4 for i in 1:P) +
        0.1 * sum((x[i+1] - x[i])^2 for i in 1:P-1))
    optimize!(model)
    @assert is_solved_and_feasible(model)
    return model, p, x
end

"Unwrap the DiffOpt.Optimizer's instantiated diff model down to the NLP Model."
function unwrap_nlp(m)
    m isa NLP.Model && return m
    m isa MOI.Bridges.LazyBridgeOptimizer && return unwrap_nlp(m.model)
    for f in (:optimizer, :model)
        hasproperty(m, f) && return unwrap_nlp(getproperty(m, f))
    end
    return error("cannot unwrap $(typeof(m))")
end

"One `reverse_differentiate!` with a fixed primal seed. The result stays in `back_grad_cache`."
function reverse_once!(model, x, Δx)
    DiffOpt.empty_input_sensitivities!(model)
    for i in eachindex(Δx)
        DiffOpt.set_reverse_variable(model, x[i], Δx[i])
    end
    DiffOpt.reverse_differentiate!(model)
    return nothing
end

"The pre-PR dense algorithm on the SAME instantiated diff model and the SAME (already-mapped)
seeds: `∂s = -K⁻¹N` (kkt_dim × P dense solve) then `Δp = ∂sᵀ Δw` — master's exact reverse
contraction (dx-only seeds). Returns the same `ci => Δp` dict shape as `ReverseCache`."
function dense_once(nlp)
    cache = nlp.cache
    ∂s, _ = NLP._compute_sensitivity(nlp)
    Δw = zeros(size(∂s, 1))
    for (i, vi) in enumerate(cache.primal_vars)
        Δw[i] = get(nlp.input_cache.dx, vi, 0.0)
    end
    Δp = ∂s' * Δw
    form = nlp.model
    return Dict(ci => Δp[form.var2param[vi].value] for (vi, ci) in form.var2ci)
end

@printf("DiffOpt %s   Julia %s   Ipopt %s\n", pkgversion(DiffOpt), VERSION, pkgversion(Ipopt))
@printf("%8s %9s | %11s %8s | %11s %8s | %7s %12s\n",
        "P", "kkt dim", "reverse ms", "GiB", "dense ms", "GiB", "ratio", "max rel Δp")
for P in (50, 200, 800, 2000)
    model, p, x = build_and_solve(P)
    Δx = [sin(0.7 * i) for i in 1:P]
    reverse_once!(model, x, Δx)                       # instantiates the diff model
    nlp = unwrap_nlp(JuMP.backend(model).diff)
    c = nlp.cache
    kkt_dim = length(c.primal_vars) + 2 * (length(c.leq_locations) + length(c.geq_locations)) +
              length(c.cons) + length(c.has_low) + length(c.has_up)

    rev = [(GC.gc(); @timed reverse_once!(model, x, Δx)) for _ in 1:REPS]
    den = [(GC.gc(); @timed dense_once(nlp)) for _ in 1:REPS]

    d_rev = nlp.back_grad_cache.Δp
    d_den = den[end].value
    scale = maximum(abs, values(d_den))
    maxrel = maximum(abs(d_rev[k] - d_den[k]) for k in keys(d_den)) / scale
    revms = median([1e3 * s.time for s in rev]); revgb = median([s.bytes / 2^30 for s in rev])
    denms = median([1e3 * s.time for s in den]); dengb = median([s.bytes / 2^30 for s in den])
    @printf("%8d %9d | %11.1f %8.3f | %11.1f %8.3f | %6.1fx %12.2e\n",
            P, kkt_dim, revms, revgb, denms, dengb, denms / revms, maxrel)
end
println("\nOn the PR branch `reverse` is the single adjoint solve; `dense` is the pre-PR")
println("algorithm (unchanged in-tree, still used by forward mode). Same Δp to ~1e-11 relative —")
println("the only difference is floating-point solve order (1 transposed RHS vs P columns).")

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.94%. Comparing base (f3fdbd8) to head (6047538).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #369      +/-   ##
==========================================
+ Coverage   91.88%   91.94%   +0.06%     
==========================================
  Files          17       17              
  Lines        2722     2745      +23     
==========================================
+ Hits         2501     2524      +23     
  Misses        221      221              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[yeabbratz]

Updated for review: applied JuliaFormatter v2 at the repo's margin = 80 to green format-check, and added test_reverse_adjoint_singular_factorization_fallback covering the K === nothing branch in _compute_sensitivity_adjoint (the line codecov flagged).

Note on the formatting commit: the files it touches are all under docs/src/examples/ and are pre-existing formatting drift on master under JuliaFormatter 2.10 (master's own format-check is currently failing for the same reason) — none of this PR's own code needed reformatting. Happy to split that into a separate cleanup PR against master if you'd prefer.