Widen UJacobianWrapper.p for nested ForwardDiff through Rosenbrock (#3381)

[ChrisRackauckas-Claude]

Summary

• Fixes ForwardDiff errors for NonlinearSolve over ODE solve #3381. When a Rosenbrock integrator is invoked with a Vector{<:Dual} p (because we are inside an outer ForwardDiff layer — e.g. a NonlinearLeastSquares parameter fit), the inner Jacobian widens u into a deeper nested-Dual type via its JacobianConfig, but the stored p in UJacobianWrapper is still at the outer Dual level. The user body then multiplies p[i] * u[i] across two different Dual nesting levels, which dispatches through ForwardDiff's tagcount-based tag precedence. That precedence is a @generated function whose literal value is baked at first compile, so its ordering depends on which package's precompile cache first instantiated each tag type — in this nested-AD scenario it inverts the nesting and produces a triple-nested Dual that crashes setindex!(du, ...) with MethodError: no method matching Float64(::Dual{..., Dual{..., Float64, 2}, 2}).
• jacobian! now inspects the prepared ForwardDiff.JacobianConfig to recover the inner xdual buffer's element type and lifts uf.p into that nested-Dual type via a fresh UJacobianWrapper before delegating to DI.jacobian!. The widened p carries zero inner partials (correct — p does not depend on u). The fast path (_widen_uf_p_for_jac(f, prep) = f) is a single type-stable dispatch so non-nested calls incur no overhead.
• Adds a regression test that builds an outer Dual p with a throwaway tag type and solves a Rosenbrock23 ODE, exercising the exact call graph that crashes pre-fix and succeeds post-fix. The test is self-contained within OrdinaryDiffEqDifferentiation's test deps (no NonlinearSolve needed).
• Bumps OrdinaryDiffEqDifferentiation patch version 2.8.0 → 2.8.1.

Root cause (abbreviated)

The user's MWE from #3381:

using OrdinaryDiffEqRosenbrock, NonlinearSolve
function ode(du, u, p, t)
    du[1] = -p[1]*u[1]
    du[2] = -u[1] - p[2]*u[2]
end
ode_f = ODEFunction{true, SciMLBase.FullSpecialize}(ode)
solve_ode(p) = solve(ODEProblem(ode_f, [1.0, 1.0], (0.0, 10.0), p),
                     Rosenbrock23(autodiff=AutoForwardDiff(chunksize=2)))
experiment = solve_ode([1.5, 2.0])
function resid!(du, p, _)
    du .= sol.(experiment.t, idxs=1) .- experiment[1, :] where sol = solve_ode(p)
end
nls = NonlinearFunction{true, SciMLBase.FullSpecialize}(resid!, resid_prototype=zeros(length(experiment.t)))
solve(NonlinearLeastSquaresProblem(nls, [1.0, 1.0]), LevenbergMarquardt())

fails with FirstAutodiffJacError wrapping
MethodError: no method matching Float64(::ForwardDiff.Dual{Tag{OrdinaryDiffEqTag, Dual{NonlinearSolveTag, Float64, 2}}, Dual{NonlinearSolveTag, Float64, 2}, 2})

The nested dual ends up being assigned into a vector whose eltype is only the outer single-level dual. Bisecting shows the broken tagcount literal order for the outer Tag{NonlinearSolveTag, Float64} vs. the inner Tag{OrdinaryDiffEqTag, Dual{NonlinearSolveTag, Float64, 2}} — the former is baked by NonlinearSolveBase's precompile workload at precompile time while the latter is generated lazily at runtime, which inverts ≺ for the cross-tag p[i]*u[i] multiplication inside ode. That inversion is what this patch defends against by normalizing p on the inner solver's side so the cross-tag arithmetic never happens.

The NonlinearSolveBase-side root cause (hardcoded Tag{NonlinearSolveTag, Float64} stamped via isa Vector{Float64} specialization) is addressed in a companion PR to SciML/NonlinearSolve.jl. This OrdinaryDiffEq PR is the robust bug-class fix: it protects any caller that hands UJacobianWrapper a Vector{<:Dual} p, regardless of how that caller's tag was registered.

Test plan

• Added lib/OrdinaryDiffEqDifferentiation/test/nested_forwarddiff_tests.jl. Reproduces the crash pre-fix (FirstAutodiffJacError on an unpatched repro env) and passes post-fix on the same env with just the derivative_wrappers.jl change developed in.
• Existing Autodiff Error Tests / Differentiation Trait Tests / No Jac Tests still pass locally (no behavior change for single-level AD — _widen_uf_p_for_jac returns f unchanged in that path).
• CI on OrdinaryDiffEqDifferentiation test group.

🤖 Generated with Claude Code

[ChrisRackauckas-Claude]

Triaged the CI failures on the initial push. Breakdown:

Caused by this PR (fixed in force-push):

• test (AD, 1.11) — test/ad/ad_tests.jl:301, @test !iszero(DI.gradient(of_a, AutoForwardDiff(), [1.0])). This is a nested ForwardDiff over an ODE solve where the inner ODEProblem(f_a, u0, tspan, p) uses the default DEFAULT_SPECIALIZATION === AutoSpecialize, so prob.f.f is a FunctionWrappersWrapper. DiffEqBase pre-compiles the FWW slot (Vector{nested_u_dual}, Vector{nested_du_dual}, Vector{outer_p_dual}, Float64) exactly for this case, and the unwidened call path matches it. My helper was widening uf.p to a nested Dual, which produced (nested_u, nested_du, *nested_p*, Float64) — a signature no FWW slot covers — which FWW v1's AllowNonIsBits policy then rejects with NoFunctionWrapperFoundError (surfacing as the FirstAutodiffJacError in the test output).

  Fix: gate _widen_uf_p_for_jac on f.f not being a FunctionWrappersWrapper. Under AutoSpecialize we leave uf.p alone and let the DiffEqBase-precompiled FWW slot take the call; under FullSpecialize (the call path from ForwardDiff errors for NonlinearSolve over ODE solve #3381 MWE and from the new regression test, where ode_f.f is a plain function) we still widen and defuse the cross-tag multiplication. Both repro_fix's MWE run and the nested_forwarddiff_tests.jl regression test still pass locally, and the minimized ad_tests.jl:301 repro now returns g = [-0.3487867031164756] instead of throwing.

Pre-existing, not caused by this PR:

• runic — drift in lib/OrdinaryDiffEqAMF/test/test_adjoint_fd2d.jl, lib/OrdinaryDiffEqCore/src/alg_utils.jl, lib/OrdinaryDiffEqNonlinearSolve/test/nested_ad_nlsolvealg_tests.jl, lib/OrdinaryDiffEqRosenbrock/test/ode_rosenbrock_tests.jl. None touched by this PR.
• test (Regression_II, *) — same 4 versions failing on current master (run 24293597671 on a6c6e192).
• test (OrdinaryDiffEqRosenbrock, 1) / test (OrdinaryDiffEqNonlinearSolve, 1) — @inferred failures of the shape return type ODESolution{...,Nothing} does not match inferred return type Union{ODESolution{...,Nothing}, ODESolution{...,JacReuseState{Float64}}}. The JacReuseState Union is inferred from the Add Jacobian reuse for Rosenbrock-W methods commit (Add Jacobian reuse for Rosenbrock-W methods #3075 / 689d0b7), not from anything in this PR.
• ModelingToolkit.jl/*, test (Integrators_II, *), test (OrdinaryDiffEqSDIRK*, *), test (StochasticDiffEq_WeakConvergence1) — from a quick pass these are either the same pre-existing inference regressions or master flakes; I have not seen this PR's diff reach any of those code paths in a local bisect. If any of them are still failing on the new run and do look PR-adjacent I'll dig in further.

Force-pushed d72f69b8a with the gated widener + the updated inline comment explaining the AutoSpecialize short-circuit.