Skip to content

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

Draft
ChrisRackauckas-Claude wants to merge 5 commits into
SciML:masterfrom
ChrisRackauckas-Claude:opt-prep-fallbacks
Draft

ChrisRackauckas-Claude wants to merge 5 commits into
SciML:masterfrom
ChrisRackauckas-Claude:opt-prep-fallbacks

Conversation

@ChrisRackauckas-Claude

@ChrisRackauckas-Claude ChrisRackauckas-Claude commented Oct 10, 2026 •

Copy link
Copy Markdown
Member

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 /opt-prep/log.txt; orchestrated by Claude Code (claude-opus-5-5[1m]) https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv

Review findings addressed

Review by codex-cli 0.162.1 / gpt-6.1-sol (CHANGES at 7a060d1); fixed in d18b4c6.

  1. High: the foreign-p Lagrangian fallback differentiated constraints at the construction p. Fixed. The in-place Lagrangian now evaluates constraints through a p-accepting term cons_lagrangian(θ, λ, p), which the out-of-place method shares. Matrix and packed-vector lag_h therefore use the call-time p on both the prepared and the prep-free path. This also fixes a same-type p ≠ p0, which returned the wrong Hessian before this PR too; Ipopt hits that case when it forwards a cache.p that reinit! replaced.
  2. Medium: σ = 0 bypassed the fallback. Fixed. In-place σ = 0 uses the prepared per-constraint Hessians only when θ has the construction type and p === the construction p, because those preps close over it. In every other case it takes the prep-free Hessian of cons_lagrangian at the call-time p; that includes a user-supplied cons_h, where no AD prep exists. Out-of-place σ = 0 used to return λ .* cons_h!(θ), which is a vector of matrices and throws a MethodError whenever cons_h = false. It now returns the Hessian of cons_lagrangian.
  3. From the earlier Devin review (non-blocking items 1 and 3): the lag_h prep gate now also checks the types of σ and λ. Before, an Int σ or a Float32/view λ threw PreparationMismatchError. The new tests compare against an independent ForwardDiff.hessian of the explicit Lagrangian.

Verification

  • New testset lag_h at the call-time p, σ = 0, and foreign σ/λ types in test/AD/dual_tolerant_tests.jl. It runs AutoForwardDiff and AutoReverseDiff over 10 cases: a call-time p of the same type and of a foreign type, a Float32 θ, σ = 0 at the default p, at another p and with a Float32 θ, σ::Int, a Float32 λ and a view λ. Each case is checked for IIP matrix, IIP packed and OOP. It also covers σ = 0 with a user cons_h.
    • Before (source change stashed): Pass 70, Fail 8, Error 3. The failures are wrong Hessians at a call-time p, MethodError: objects of type Nothing are not callable at OOP σ = 0, and PreparationMismatchError at σ = 0 with a Float32 θ. The testset aborts at the first uncaught error.
    • After: dual_tolerant_tests.jl 126/126 on Julia 1.13.2, and 126/126 on 1.10.12 (LTS).
  • OPTIMIZATION_TEST_GROUP=AD (1.13.2): 1106/1106. Core: 73/73. QA: 20 pass, 1 broken (pre-existing).
  • Matching-type lag_h (IIP σ = 1, IIP σ = 0, IIP packed, OOP), checked under ForwardDiff and ReverseDiff. Allocations are identical under ForwardDiff and within 400 B under ReverseDiff; timings are unchanged.
  • Probe under AutoEnzyme and AutoReverseDiff(compile = true): IIP, packed and OOP lag_h at σ = 1 and σ = 0 match the reference. Runic and typos report nothing.

Still open (pre-existing, not changed here)

  • AutoSymbolics with lag_h = true fails at instantiation, with and without this commit. The error is MethodError: erase_cache_vars! from DI's Symbolics extension, raised while preparing the Lagrangian Hessian.
  • AutoFiniteDiff second-order results are inaccurate. DI's own hessian(f, SecondOrder(AutoFiniteDiff(), AutoFiniteDiff()), x) gives [24 -12; -12 8] where the exact Hessian is [23.24 -8.4; -8.4 6], and OptimizationBase maps AutoFiniteDiff to that SecondOrder.
  • The in-place Lagrangians in OptimizationZygoteExt and OptimizationDISparseExt still use cons_oop(θ), which closes over the construction p. They have the same call-time-p bug.
  • The σ = 0 fallback is prep-free and re-prepares on every call. This applies to in-place σ = 0 with a non-construction p or θ type, and to every out-of-place σ = 0 call.

🤖 Posted by an AI agent — harness: Claude Code 2.1.301 · model: claude-opus-5-5
Conversation: local session ID 71155362-c106-40f7-b5a3-f1cf335f4339 (headless, no URL)

Risk assessment (independent review, codex-cli gpt-6.1-sol, reviewed head d18b4c6)

  • Risk: medium
  • Blast radius: Dense DifferentiationInterface-generated gradients, value/gradient/Hessian combinations, Hessian-vector products, constraint Jacobians, and Lagrangian Hessians, in both in-place and out-of-place forms. The helpers sit on common AD paths. No public name/signature is removed, and no runtime dependency is added.
  • Evidence: Reviewed head d18b4c6b3af4e5900f2169c3daa083c715ac0bf4, unchanged at the final check. Julia 1.13.0, four threads: regression file 126/126, official Core group 73/73, independent polynomial probe 738/738, and the earlier review's exact reproductions 2/2. The current regression file on pre-fix head 7a060d15f0d8c5573873ee15f5ffbd828a795c74 produces 70 pass, 8 fail, 3 error. All 14 CI checks were pending at both the initial and final checks; there are no failed PR jobs/logs to compare yet. CI: https://github.com/SciML/Optimization.jl/actions/runs/38924793086 and https://github.com/SciML/Optimization.jl/actions/runs/38924793437 . Documentation: https://github.com/SciML/Optimization.jl/actions/runs/38924792563 . Downstream: https://github.com/SciML/Optimization.jl/actions/runs/38924793325 . No CI-green claim is made.
  • Independent review: codex-cli 0.163.0 gpt-6.1-sol rated it medium, high confidence: both earlier defects discriminate before/after, all targeted regressions pass, and independent numerical checks cover ordinary and compiled ReverseDiff plus ForwardDiff; broader backend and downstream coverage remains pending.
  • Merge: needs human review — medium blast radius across derivative dispatch and entirely pending CI prevent an auto-merge recommendation. No blocking correctness finding remains in the reviewed change.

🤖 Posted by an AI agent — harness: codex-cli · model: gpt-6.1-sol (independent reviewer; appended by Claude Code 2.1.301 · claude-opus-5-5)
Conversation: https://claude.ai/code/session_01VZjwf7Ux9YtTzwHGHMqnsq

ChrisRackauckas and others added 3 commits October 10, 2026 02:30
Extend the existing `_prep_valid` gate beyond grad!/cons_j! so second-order
and combined DI operators fall back to prep-free calls when θ/p eltypes or
array types differ from construction, avoiding PreparationMismatchError.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Cursor Agent <noreply@cursor.com>
Agent-Harness: Cursor Agent 2026.10.01-14929f9
Agent-Model: unknown (Cursor auto)
Agent-Session: local session, transcript /home/crackauc/sandbox/goals/performance/jobs/opt/opt-prep/log.txt on amdci2
Share Type-param helper closures (incl. grad!/cons_j!/lag_h!), drop
duplicated NullParameters branches, use value_gradient_and_hessian(!)
for fgh!, and gate hv! on the tangent type.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Cursor Agent <noreply@cursor.com>
Agent-Harness: Cursor Agent 2026.10.01-14929f9
Agent-Model: unknown (Cursor auto)
Agent-Session: local session, transcript /home/crackauc/sandbox/goals/performance/jobs/opt/opt-prep/log.txt on amdci2
Named _iip/_oop_ad_cons_jac_clos match the other Type-param factories,
build hv prep tangents from zero(x) so SVector v hits the prepared path,
and cover IIP packed-vector plus OOP lag_h foreign types.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Cursor Agent <noreply@cursor.com>
Agent-Harness: Cursor Agent 2026.10.01-14929f9
Agent-Model: unknown (Cursor auto)
Agent-Session: local session, transcript /home/crackauc/sandbox/goals/performance/jobs/opt/opt-prep/log.txt on amdci2
@ChrisRackauckas

Copy link
Copy Markdown
Member

🤖 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


🤖 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)

ChrisRackauckas and others added 2 commits October 16, 2026 06:36
The in-place Lagrangian differentiated `cons_oop(θ)`, which closes over the
construction `p`, so `lag_h(H, θ, σ, λ, p)` with any other `p` returned the
constraint term at the construction parameters. With the prep-free fallback
accepting a foreign `p` type, that wrong result was returned silently instead
of a PreparationMismatchError. Both Lagrangians now share a `p`-accepting
constraint term `cons_lagrangian(θ, λ, p)`.

σ = 0: the in-place path uses the prepared per-constraint Hessians only at
the construction θ type and for the construction `p` object (they close over
it). Otherwise, and when no AD constraint-Hessian preparation exists (a
user-supplied `cons_h`), it takes the prep-free Hessian of `cons_lagrangian`
at the call-time `p`. The out-of-place σ = 0 path returned `λ .* cons_h!(θ)`
(a vector of matrices, and a MethodError when `cons_h = false`); it now
returns the Hessian of `cons_lagrangian`.

The Lagrangian preparation is built at `Constant(σ0)`/`Constant(λ0)`; their
types join the prep gate, so an `Int` σ or a `Float32`/view λ takes the
prep-free path instead of throwing PreparationMismatchError.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Agent-Harness: Claude Code 2.1.301
Agent-Model: claude-opus-5-5
Agent-Session: local session ID 71155362-c106-40f7-b5a3-f1cf335f4339 (headless claude -p, no URL)

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants