[AI] OptimizationBase: use a supplied f.grad when synthesizing fg!

[ChrisRackauckas-Claude]

Please ignore until reviewed by @ChrisRackauckas. Opened as a draft.

Fixes #1282.

Root cause

Every AD extension gates fg! construction on fg == true && f.fg === nothing — it never checks f.grad. So when a caller supplies a gradient but not an fg, OptimizationBase builds a brand-new AD preparation over f.f and the supplied gradient is never called.

That was harmless while allowsfg was off for Optim. #1200 flipped it on, so every first-order line search started taking the rebuilt path.

Pathfinder.jl hits this squarely: it supplies grad wrapping LogDensityProblems.logdensity_and_gradient (Pathfinder/src/optimize.jl:26-28), i.e. Turing's own prepared AD path. Verified with a counter inside a user-supplied gradient over 100 LBFGS iterations:

allowsfg=true : user grad calls = 0,   raw primal calls = 0
allowsfg=false: user grad calls = 151, raw primal calls = 151

The replacement is far slower because Pathfinder's objective closure captures the whole LogDensityProblem (model + data + VarInfo), so Mooncake's tangent is that entire object graph and has to be zeroed before every reverse pass.

Measurements

66-parameter hierarchical Turing model, 2400 observations, Mooncake, Optim.LBFGS(), Julia 1.11:

	ms/call
Turing logdensity_and_gradient (what Pathfinder supplied)	0.372
OptimizationBase-rebuilt DI value_and_gradient!	10.269

Flat profile of the rebuilt path: 2948 of 3085 samples (96%) in Mooncake.set_to_zero!! / set_to_zero_internal!! / _already_tracked!. Only ~4% is actual derivative work.

End to end:

	allowsfg=false (pre-#1200)	allowsfg=true (master)	this PR
pathfinder 1st	23.6 s	49.4 s	24.0 s
pathfinder 2nd	0.55 s	15.3 s	0.53 s
multipathfinder(8)	18.1 s	139.2 s	19.8 s

The first-run component is the extra prepare_gradient — 70 s cold, 12.4 s warm — paid once per solve, i.e. once per Pathfinder path.

Change

Prefer a supplied f.grad when synthesizing fg!, in OptimizationDIExt.jl (in-place and out-of-place), OptimizationZygoteExt.jl (dense and sparse), and OptimizationEnzymeExt.jl (both):

fg! = if fg == true && f.fg === nothing && f.grad !== nothing
    (res, θ, p = p) -> (f.grad(res, θ, p); f.f(θ, p))
elseif fg == true && f.fg === nothing
    ... existing AD path ...

Also drops the duplicated prepare_gradient in the DI extension, whose two branches were identical despite the comment claiming one reused the gradient preparation.

This is not Optim-specific: OptimizationOptimisers, OptimizationSophia, and OptimizationAuglag also set allowsfg = true and have been discarding user gradients the same way. They are covered too.

OptimizationBase 5.2.2 → 5.2.3.

Test plan

New testset in lib/OptimizationBase/test/AD/adtests.jl covering ForwardDiff / ReverseDiff / Zygote / Enzyme / AutoSparse(AutoZygote()) × g ∈ (false, true), asserting the supplied gradient is called exactly once and that the AD path still builds its own preparation when no gradient is supplied.

Run locally on Julia 1.11:

• OPTIMIZATION_TEST_GROUP=AD Pkg.test() for OptimizationBase — 775 pass, 0 fail
• Pkg.test() for OptimizationOptimJL — 6969 pass, 0 fail
• Runic clean
• CI green

A follow-up PR will fix two unrelated pre-existing bugs in OptimizationEnzymeExt.jl's out-of-place instantiate_function noticed while working here.

🤖 Generated with Claude Code

https://claude.ai/code/session_014Uj4Pu1LKHYqLSTLfPYDPo