Fix DI prep type-mismatch fallbacks for hess/hv/fg/fgh/lag_h

[ChrisRackauckas-Claude]

Summary

Collapses the DI prep-validity fallbacks in OptimizationDIExt.jl onto shared Type-parameter helper closures (including grad!/cons_j!/lag_h!/_ad_cons_jac), fixing the duplicated NullParameters branches. Also switches AD-generated fgh! from DI's scalar value_derivative_and_second_derivative(!) to value_gradient_and_hessian(!) so the matching-type fast path works, gates hv! on the tangent v type (built from zero(x) so SVector tangents hit the prepared path), and adds foreign-type coverage for IIP packed-vector and OOP lag_h.

Out of scope (still broken, not addressed here): Dual p with second-order ops raises DualMismatchError (ForwardDiff tag nesting, even with a fresh Dual-p prep); OOP SVector solves through OptimJL/Optimisers remain wrapper bugs beyond the unit-level instantiate_function fallback.

Verification

Regression (dual_tolerant_tests.jl):

• before (src stashed): matching-type fgh! → MethodError on value_derivative_and_second_derivative!
• after: matching-type fgh and foreign-type hess/hv/fg/fgh/lag_h (incl. IIP packed-vector + OOP lag_h) pass

Matching-type path unchanged (n=10, min ns, master vs PR; identical allocations):

op	master	PR
iip grad!	137.7	139.1
iip hess!	2528	2553
iip fg!	147.4	146.8
iip hv!	561	565
oop hess	2797	2751
oop fg	170.3	169.4
oop hv	604	606

Test groups (Julia 1, OPTIMIZATION_TEST_GROUP):

• AD: 1045 Pass
• QA: 20 Pass + 1 Broken

Probes: Float32/BigFloat/SVector/hv-with-mismatched-v and matching fgh! match fresh-prep results; Dual-p second-order still DualMismatchError (out of scope).

Not verified

• Full end-to-end a4_generic.jl Optim/Optimisers SVector solve matrix (unit-level covered; wrapper bugs out of scope).
• OptimizationDISparseExt second-order gates (still only cons_j-style).

Reviewer pushback

• Whether Dual-p second-order / Optim SVector wrappers should be in-scope follow-ups.
• Whether the Type-param helper surface is the right long-term shape vs a macro.

Please ignore until reviewed by @ChrisRackauckas.

Risk assessment

Independent review: Claude Code (claude-opus-5-5), 3 rounds; round 3 MERGE, risk low (confidence medium-high): matching-type path identical to master (prepared calls, allocations, timings), 46 foreign-type probe cases match a fresh prep, AD 1045/1045 and QA pass; remaining CI failures are known on master (MTK downstream, Reactant downgrade, Manopt Core) or Python-setup infrastructure (PyCMA, SciPy).

🤖 Generated with Cursor Agent 2026.10.01-14929f9 (model: unknown (Cursor auto)), transcript /home/crackauc/sandbox/goals/performance/jobs/opt/opt-prep/log.txt on amdci2; orchestrated by Claude Code (claude-opus-5-5[1m]) https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv

[ChrisRackauckas]

🤖 Automated comment from an AI agent running as @ChrisRackauckas — not written or reviewed by Chris.

Independent review (Devin CLI 3000.11.3, model fusion-claude-opus-5-5-high-sidekick-swe-2-medium): MERGE, risk medium.

Full review

VERDICT: MERGE
RISK: medium

The change sits on the hot path of every DI-backed instantiate_function (AutoForwardDiff / AutoReverseDiff / AutoFiniteDiff / SecondOrder, both IIP and OOP). It changes no public API, and the matching-type behaviour stays the same. That hot-path exposure is why the risk is medium and not low.

Blocking findings

None.

What I confirmed by running (details under "What I ran"):

• The headline bug is real and this PR fixes it. On master, the AD-generated fgh throws a MethodError for every DI backend I tried (ForwardDiff, ReverseDiff, FiniteDiff; IIP and OOP). Master calls value_derivative_and_second_derivative(!) with the prep where DI expects the backend, and the OOP branch also passes adtype instead of soadtype. On the PR, fgh matches closed-form gradients and Hessians at p0 and at an explicit p1.
• The new fallbacks work. Master throws PreparationMismatchError on IIP hess with a Float32 θ, on hv with a SubArray tangent, on lag_h with a Float32 θ, and on OOP hess/hv/fgh built at an SVector and called with a Vector. The PR gives correct results in all four cases, checked against closed-form references, not just against a fresh prep.
• Fail-before / pass-after is real. The PR's test file errors on master in matching-type fgh! (IIP/OOP) (MethodError at dual_tolerant_tests.jl:192) and passes on the PR: 65/65 across 7 testsets. Master's own dual_tolerant_tests.jl also passes on the PR source.
• The matching-type path does not regress. @allocated is identical to master for IIP hess/fg/hv and OOP hess; IIP grad allocates less (1152 vs 1456 B). OOP grad is now @inferred-clean; on master it inferred Any. The OOP hess/hv/fg/fgh results are @inferred-clean.

Non-blocking findings

1. lag_h gate ignores the σ and λ types (OptimizationDIExt.jl:149-213, _iip_lag_h_clos / _oop_lag_h_clos). Confirmed by running. _lag_prep is built with Constant(one(eltype(x))) and Constant(ones(eltype(x), num_cons)), so DI also checks the context types. With a matching Float64 θ, all of these still throw PreparationMismatchError (contexts: ❌) on the PR:

   • σ::Float32
   • σ::Int
   • λ::Vector{Float32}
   • λ passed as a view, in the IIP matrix form, the IIP packed form and the OOP form

   This holds for ForwardDiff, ReverseDiff and FiniteDiff. Master behaves the same, so it is not a regression. But the PR advertises "foreign-type … lag_h" coverage, and this is the same class of bug in the same closure. MOI and Ipopt dodge it with Vector(μ). OptimizationNLPModels.jl:258/261 passes the solver's y straight through. Suggested fix: add _prep_valid(Tσ0, σ) && _prep_valid(Tλ0, λ) to the gate, with Tσ0 = typeof(one(eltype(x))) and Tλ0 = typeof(ones(eltype(x), num_cons)), plus a test.

2. The σ = 0 branch of IIP lag_h is still ungated. Confirmed by running. With a Float32 θ and σ = 0, the call goes to cons_h_weighted!, which uses prep_cons_hess[i] with no fallback, and throws PreparationMismatchError (x: ❌). cons_h (IIP), and cons_vjp with ReverseDiff, throw the same way for a Float32 θ. The PR body only lists the SparseExt gates as unverified; it should also name these as still broken.

3. The new lag_h tests have no independent reference (dual_tolerant_tests.jl:290-314). Every lag_h assertion compares the gated closure with a fresh prep, i.e. the same DI code path. I checked by hand against σ∇²f + λ∇²c in closed form, and the values are correct for ForwardDiff and ReverseDiff. The test should still assert against a closed-form or ForwardDiff.hessian reference, as the hess/hv/fgh cases do.

4. The Float32 tolerances (rtol = 1.0e-3) are wider than the precision justifies. In my probes the Float32 results agree with closed form to rtol = 1e-5, so 1e-4 would still pass with margin and catch more. This is minor and not fitted to observed error in a way that hides anything.

5. Pre-existing bugs I found while probing. None are caused by this PR. Each deserves its own issue, not scope creep here:

   • IIP lag_h(H, θ, σ, λ, p1) with an explicit p ≠ p0 returns the wrong Hessian on master and on the PR. The reason is that lagrangian (OptimizationDIExt.jl:409-411) calls cons_oop(θ), which bakes in the construction p, so the constraint term ignores p1. OOP lag_h is correct. This overlaps with Keep instantiated closures in sync with the live p (#1383) #1400 (live p).
   • OOP lag_h at σ = 0 throws MethodError: objects of type Nothing are not callable even for a matching Float64 θ, because cons_h! is nothing unless cons_h = true (OptimizationDIExt.jl:~805).
   • AutoFiniteDiff Hessians are about 15% off. OptimizationBase gives [20 -12; -12 6] against an exact [23.24 -8.4; -8.4 6]. DI's own hessian(f, SecondOrder(AutoFiniteDiff(), AutoFiniteDiff()), x, Constant(p)) returns [24 -12; -12 8], while plain AutoFiniteDiff() is correct. The problem is in generate_adtype's FD-over-FD choice and/or upstream, not in this PR.

6. Conflict with sibling work. OptimizationBase: make implicit imports explicit, un-break no_implicit_imports QA check #1391 edits the import block of OptimizationDIExt.jl, and its hunk context contains the value_derivative_and_second_derivative line this PR removes. Expect a trivial rebase conflict. I don't expect a textual conflict with Keep instantiated closures in sync with the live p (#1383) #1400, since it doesn't touch this file, but its instantiate_live wraps these closures. Neither PR has merged yet.

7. Scope and style. The PR combines a refactor (collapsing the duplicated NullParameters branches into typed factories) with three fixes: the fgh MethodError, the second-order fallbacks, and the hv tangent gating. These are coupled closely enough to accept as one PR. Comment density is low. The test comment "Matching-type fgh uses value_gradient_and_hessian(!)" describes the implementation rather than the behaviour under test and could go. There is no linked issue or tracking issue, and no OptimizationBase version bump (the fix is a patch). The only new dependency is StaticArrays (MIT), and only in the test environment.

What I ran

All runs used Julia 1.11 with DI 0.7.21, ForwardDiff 1.4.6, ReverseDiff 1.18.5, FiniteDiff 2.33.0, Enzyme 0.13.216 and SciMLBase 3.57.0. I used two scratch envs that differ only in the OptimizationBase source: PR head, or an origin/master worktree at e03638e.

• PR test/AD/dual_tolerant_tests.jl on PR source: exit 0, all 7 testsets pass (65 tests).

• PR test file on master source: errors in matching-type fgh! (IIP/OOP) with MethodError … value_derivative_and_second_derivative!, and the script aborts there.

• Master's dual_tolerant_tests.jl on PR source: exit 0, 5 testsets pass.

• scratch/probe.jl, my own probe, run under both envs. It checks against closed-form gradient, Hessian, constraint-Hessian and Lagrangian references for AutoForwardDiff, AutoReverseDiff and AutoFiniteDiff:

  • matching types at p0 and p1
  • Float32 θ
  • a view tangent
  • foreign σ and λ
  • σ = 0 with a Float32 θ
  • cons_h / cons_vjp with a Float32 θ
  • an OOP closure built at an SVector and called with a Vector
  • @inferred on the OOP closures
  • @allocated on the matching path

  Logs are in scratch/probe_pr.log and scratch/probe_master.log; the results are summarised in the findings above.

• scratch/probe2.jl plus a direct DI call to isolate the FiniteDiff Hessian error.

• Runic --check and typos on both changed files: clean.

• gh pr view 1408, the PR checks (AD, Core and QA jobs green; the red jobs are ModelingToolkit downstream, the Reactant downgrade, Manopt Core, and PyCMA/SciPy QA), and the diffs of sibling PRs OptimizationBase: make implicit imports explicit, un-break no_implicit_imports QA check #1391 and Keep instantiated closures in sync with the live p (#1383) #1400.

What I did not verify

• The full OPTIMIZATION_TEST_GROUP=AD (claimed 1045 pass) and QA groups locally. I skipped them because of host limits. CI shows both green on this head.
• The PR body's timing table. I measured allocations only, not ns timings.
• The PR body's claim that the CI failures already exist on master (MTK downstream, Reactant downgrade, Manopt Core, PyCMA/SciPy). I did not open those logs.
• The Dual-p second-order DualMismatchError that the PR body declares out of scope, and the Optim/Optimisers SVector end-to-end solves.
• AutoZygote's DI second-order path (SecondOrder(AutoForwardDiff, AutoZygote)) and the Mooncake path through this file.
• The OptimizationDISparseExt gates.

Links

• PR under review: Fix DI prep type-mismatch fallbacks for hess/hv/fg/fgh/lag_h #1408
• Sibling PR (import-block conflict): OptimizationBase: make implicit imports explicit, un-break no_implicit_imports QA check #1391
• Sibling PR (live p, related to the IIP lag_h explicit-p bug): Keep instantiated closures in sync with the live p (#1383) #1400

---

🤖 Posted by an AI agent — harness: Devin CLI 3000.11.3 (review), Claude Code 2.1.285 (posting) · model: fusion-claude-opus-5-5-high-sidekick-swe-2-medium
Conversation: local Claude session 3c311569-dda6-4687-bc9f-65febb212c33 (fleet master)