Avoid wide Enzyme batches for Lagrangian Hessians

[ChrisRackauckas-Claude]

Ignore this PR until it has been reviewed by @ChrisRackauckas.

What changed

Compute Enzyme-generated Lagrangian Hessians one direction at a time instead of compiling one forward-over-reverse batch whose width equals the optimization dimension. The previous batch has superlinear LLVM/Enzyme compilation growth and prevents the OptimizationFrameworks/clnlbeam.jmd benchmark from reaching its solver comparisons.

The row-wise sweep zeroes and differentiates a reusable gradient workspace through a Lagrangian closure. The dense entry point copies the computed triangle symmetrically; the packed entry point consumes the same directional results. The regression test checks both representations, including the objective's off-diagonal term, for optimization dimensions through 183 variables.

The final head 91ee2925cfcc01b1389955542ec74e6f1b8f599e is rebased onto current Optimization master c2d6a1d45d8b6367620de199669fc57d4a838784, after the width-8 batching change and OptimizationBase 5.5.2 release.

Failing before / passing after

Original wide-batch failure

I ran the committed regression test against an unfixed checkout and this branch:

timeout 3600 /usr/bin/time -f 'wall=%e maxrss_kb=%M exit=%x' \
  ~/.juliaup/bin/julia +1.11.9 --startup-file=no \
  --project=lib/OptimizationBase/test/AD \
  lib/OptimizationBase/test/AD/enzyme_lagrangian_hessian.jl

Before, the command exhausted the complete one-hour Bash timeout (exit 124) while compiling the width-183 Enzyme batch. The termination stack was in Enzyme.removeDeadArgs! / LLVM Attributor, called from the old lag_h!.

On the rebased final head, the same command completes:

Test Summary:             | Pass  Total   Time
Enzyme Lagrangian Hessian |   26     26  29.6s

Julia 1.12 regression in the first row-wise implementation

The first pushed row-wise implementation (53ebab8ac4edb990e1f8b50ae5f1c43c4ca52634) failed the same test locally on Julia 1.12.7 with NaN Hessian entries at N = 10, 20, and 60:

timeout 3600 ~/.juliaup/bin/julia +1.12.7 --compiled-modules=no \
  --project=lib/OptimizationBase/test/AD \
  -e 'include("lib/OptimizationBase/test/AD/enzyme_lagrangian_hessian.jl")'

Test Summary:             | Pass  Fail  Total
Enzyme Lagrangian Hessian |   23     3     26
ERROR: LoadError: Some tests did not pass: 23 passed, 3 failed, 0 errored, 0 broken.

The exact Julia 1.12 command passes on the rebased final head:

Test Summary:             | Pass  Total     Time
Enzyme Lagrangian Hessian |   26     26  1m04.2s

--compiled-modules=no avoids an unrelated RuntimeGeneratedFunctions world-age precompile failure in the isolated Julia 1.12 depot.

Benchmark verification

I developed each OptimizationBase checkout into the SciMLBenchmarks OptimizationFrameworks environment and ran:

timeout 3600 ~/.juliaup/bin/julia +1.11.9 --startup-file=no \
  --project=. benchmark.jl benchmarks/OptimizationFrameworks/clnlbeam.jmd

The unfixed benchmark exhausted the one-hour timeout while starting run_enzyme_diff(60). With this branch it completed every chunk through size 200 in 11m55.7s and generated both clnlbeam.md and clnlbeam_7_1.png.

Test verification

GROUP=OptimizationBase_AD JULIA_PKG_PRECOMPILE_AUTO=0 \
  timeout 3600 ~/.juliaup/bin/julia +1.11.9 --project=. \
  -e 'using Pkg; Pkg.test()'

Test Summary: | Pass  Total      Time
AD            |  827    827  15m25.7s
Testing OptimizationBase tests passed
Testing Optimization tests passed

GROUP=QA JULIA_PKG_PRECOMPILE_AUTO=0 \
  timeout 3600 ~/.juliaup/bin/julia +1.11.9 --project=. \
  -e 'using Pkg; Pkg.test()'

Test Summary: | Pass  Total     Time
QA            |   21     21  1m26.0s
Testing Optimization tests passed

Runic, typos, and git diff --check passed on the final three-file diff.

Not verified locally

I did not run every test group in the Optimization monorepo. The changed extension is covered by the complete OptimizationBase AD group and root QA group. This changes no public API or documentation.

🤖 Generated with Claude Code

https://chatgpt.com/codex/tasks/01a03a17-ad6f-7131-82fc-d0fd57ea6512

[ChrisRackauckas-Claude]

Follow-up 11591555ed3ee2b7676bf578ece89287ca8e87f1 fixes the Julia 1.12 NaN failures seen in the first row-wise implementation. The exact Julia 1.12 focused test passed twice (26/26 each); the full Julia 1.11 OptimizationBase AD group passed 821/821, root QA passed 21/21, and Runic/typos/diff checks are clean. The PR body now contains the exact failing-before and passing-after evidence.

[ChrisRackauckas-Claude]

CI audit after #1331 merged:

• The PR head 11591555ed3ee2b7676bf578ece89287ca8e87f1 merges cleanly into current master 9bfa701305b0a5e2913f17d1a195a2ae9c1e5e15.
• I tested that exact local merge with Julia 1.11.9:

timeout 3600 /usr/bin/time -f "wall=%e maxrss_kb=%M exit=%x" \
  ~/.juliaup/bin/julia +1.11 --startup-file=no \
  --project=lib/OptimizationBase/test/AD \
  -e 'using Pkg; Pkg.instantiate(); include("lib/OptimizationBase/test/AD/enzyme_lagrangian_hessian.jl")'

Test Summary:             | Pass  Total   Time
Enzyme Lagrangian Hessian |   26     26  30.9s
wall=71.13 maxrss_kb=976492 exit=0

• All nine red check names are identical to the red checks on merged Cap AutoEnzyme Hessian batches at width 8 #1331. Their logs fail in dependency resolution or unrelated downgraded/downstream environments before exercising this extension: docs (SimpleOptimization/SimpleNonlinearSolve), downstream (LogExpFunctions, OrderedCollections, or Tullio), and downgrade jobs (ModelingToolkitBase, ForwardDiff, or NOMAD result fields).
• The relevant hosted OptimizationBase AD tests pass on Julia 1 and LTS, as do Core and QA. Current master also contains Fix downgrade compat floors for Ipopt, NLopt, and NOMAD sublibraries #1333, which addresses three of the downgrade compat failures, so a branch update/re-run should remove stale failures from this head.

I found no PR-specific failing check. The PR remains draft and still requires review by @ChrisRackauckas.

[ChrisRackauckas]

This is too extreme: it completely removes all batch duplicated. Just like the hessians and such, shouldn't we limit it to like 8?

[ChrisRackauckas-Claude]

Addressed the review comment: Lagrangian Hessians now use the same width-8 Enzyme FoR batching as objective Hessians (_hessian_batch_width / #1331), instead of one seed at a time.

Julia 1.12 still falls back to row-wise Duplicated seeds via _lag_hessian_batch_width, because BatchDuplicated through the Lagrangian currently produces NaN Hessian rows there. On Julia ≤1.11 the batch width is capped at 8 as requested.

Verified: Enzyme Lagrangian Hessian 26/26 on Julia 1.11.9 (width 8) and Julia 1.12.7 (row-wise fallback).

[ChrisRackauckas-Claude]

Follow-up addressing the sol review:

• Fixed OOP lag_h to use lagrangian_oop (cons(x,p)) and return a dense symmetric matrix (matching OOP hess).
• Narrowed the row-wise fallback to Julia 1.12 ≤ VERSION < 1.13 with a TODO to re-check on 1.13+.
• Expanded tests: nonzero clnlbeam points (constraint Hessians contribute), quadratic σ/μ cases including σ=0, and OOP coverage at n=9 and n=17.

Enzyme Lagrangian Hessian: 82/82 on Julia 1.11.9 and 1.12.7.