Widen UJacobianWrapper.p for nested ForwardDiff (#3381)

When a Rosenbrock integrator runs inside an outer ForwardDiff layer, the
stored `p` in `UJacobianWrapper` is a `Vector{<:Dual}` at the outer Dual
level, but the inner Jacobian computation widens `u` into a deeper
nested-Dual type via its prepared `JacobianConfig`. The user
`f(du, u, p, t)` 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 a given
tag type — in nested scenarios this can invert nesting and produce a
triple-nested `Dual` that crashes `setindex!(du, ...)` with
`Float64(::nested_dual)`.

Fix: in `jacobian!`, inspect the prepared `ForwardDiff.JacobianConfig`
for the inner xdual buffer's element type and lift `uf.p` into that
nested-Dual type ahead of time via a fresh `UJacobianWrapper`. The
widened `p` carries zero inner partials (correct — `p` does not depend
on `u`), and the fast path is a single type-stable dispatch returning
`f` unchanged so there is no overhead for non-nested calls.

Add a regression test exercising the precise call graph
(outer-Dual `p`, inner Rosenbrock Jacobian) that crashes pre-fix with a
`FirstAutodiffJacError(MethodError(Float64, (nested_dual,)))` and
succeeds post-fix.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ChrisRackauckas

lib/OrdinaryDiffEqDifferentiation/Project.toml

@@ -1,5 +1,5 @@
 name = "OrdinaryDiffEqDifferentiation"
-version = "2.8.0"
+version = "2.8.1"
 uuid = "4302a76b-040a-498a-8c04-15b101fed76b"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>", "Yingbo Ma <mayingbo5@gmail.com>"]
 

lib/OrdinaryDiffEqDifferentiation/src/derivative_wrappers.jl

@@ -185,6 +185,45 @@ function jacobian(f::F, x, integrator) where {F}
     return jac
 end
 
+# Inner-Dual eltype that the prepared ForwardDiff JacobianConfig will allocate
+# for its xdual buffer. Returns `nothing` if `prep` is not a ForwardDiff
+# JacobianConfig-backed prep (e.g. AutoFiniteDiff path).
+function _jac_prep_inner_dual_eltype(prep)
+    hasfield(typeof(prep), :config) || return nothing
+    cfg = getfield(prep, :config)
+    cfg isa ForwardDiff.JacobianConfig || return nothing
+    duals = cfg.duals
+    duals isa Tuple && length(duals) >= 2 || return nothing
+    return eltype(duals[2])
+end
+
+# When the integrator's stored `p` (held in `f::UJacobianWrapper`) is a
+# `Vector{<:Dual}` because we are *inside* an outer ForwardDiff Jacobian /
+# gradient, the inner Rosenbrock Jacobian widens `u` into a deeper nested-Dual
+# type via the prepared `JacobianConfig`. The user `ode(du, u, p, t)` body
+# then multiplies `p[i] * u[i]` across two different Dual nesting levels,
+# which falls back to ForwardDiff's tag-precedence machinery (`tagcount`-
+# based `≺`). That precedence is unstable across precompile boundaries (the
+# `@generated tagcount` literal is baked at first compile and depends on
+# which package precompiled which tag first), so the result type can come
+# out wrong and crash inside `setindex!(du, ...)` with `Float64(::nested_dual)`.
+#
+# Lift `p` into the inner nested-Dual type so the user body never multiplies
+# across tag levels. The widened `p` carries zero inner partials (which is
+# what we want — `p` does not depend on `u`).
+_widen_uf_p_for_jac(f, prep) = f
+function _widen_uf_p_for_jac(f::UJacobianWrapper, prep)
+    inner_T = _jac_prep_inner_dual_eltype(prep)
+    inner_T === nothing && return f
+    p = f.p
+    p isa AbstractArray || return f
+    Tp = eltype(p)
+    Tp <: ForwardDiff.Dual || return f
+    Tp === inner_T && return f
+    inner_T <: ForwardDiff.Dual || return f
+    return UJacobianWrapper{isinplace(f)}(f.f, f.t, convert.(inner_T, p))
+end
+
 function jacobian!(
         J::AbstractMatrix{<:Number}, f::F, x::AbstractArray{<:Number},
         fx::AbstractArray{<:Number}, integrator::SciMLBase.DEIntegrator,
@@ -240,14 +279,16 @@ function jacobian!(
         config = jac_config[1]
     end
 
+    f_eff = _widen_uf_p_for_jac(f, config)
+
     if integrator.iter == 1
         try
-            DI.jacobian!(f, fx, J, config, gpu_safe_autodiff(alg_autodiff(alg), x), x)
+            DI.jacobian!(f_eff, fx, J, config, gpu_safe_autodiff(alg_autodiff(alg), x), x)
         catch e
             throw(FirstAutodiffJacError(e))
         end
     else
-        DI.jacobian!(f, fx, J, config, gpu_safe_autodiff(alg_autodiff(alg), x), x)
+        DI.jacobian!(f_eff, fx, J, config, gpu_safe_autodiff(alg_autodiff(alg), x), x)
     end
 
     return nothing

lib/OrdinaryDiffEqDifferentiation/test/nested_forwarddiff_tests.jl

@@ -0,0 +1,61 @@
+using Test
+using OrdinaryDiffEqRosenbrock
+using SciMLBase
+using ADTypes
+using ForwardDiff
+
+# Regression test for nested ForwardDiff over an ODE solve (#3381). When a
+# Rosenbrock solver is invoked with a `Vector{<:Dual}` `p` (i.e. we are
+# inside some *outer* ForwardDiff layer), the inner Rosenbrock Jacobian
+# widens `u` to a deeper nested-Dual type via its JacobianConfig. The user
+# `f(du, u, p, t)` body then multiplies `p[i] * u[i]` across two different
+# Dual nesting levels. ForwardDiff's cross-tag multiplication uses
+# `tagcount`-based precedence, which is a `@generated` function whose
+# literal value is baked at first compile. Depending on the order in which
+# tag types are first instantiated (which varies with precompile order),
+# the resulting ordering can invert nesting and produce a triple-nested
+# `Dual` that crashes the eventual `setindex!` with
+#     `Float64(::nested_dual)` MethodError.
+#
+# The fix (`_widen_uf_p_for_jac` in derivative_wrappers.jl) lifts
+# `uf.p` into the inner nested-Dual type ahead of time so the user
+# body never has to cross tag levels.
+#
+# To reliably reproduce the broken precedence without depending on
+# NonlinearSolve's precompile-baked NonlinearSolveTag, we construct the
+# outer Dual manually with a throwaway tag that is *not* a concrete
+# `Tag(f, V)` call. Because the throwaway tag is only used as a type
+# parameter, ForwardDiff does not trigger its `tagcount` until the first
+# cross-tag multiplication — which happens *after* `OrdinaryDiffEqTag` has
+# already had its nested-V `tagcount` triggered via
+# `OrdinaryDiffEqDifferentiation.prepare_ADType`. This matches the
+# precedence inversion that the original issue hits via
+# `NonlinearSolveTag`'s precompile workload.
+
+struct _Nested3381OuterTag end
+
+function _nested3381_ode!(du, u, p, t)
+    du[1] = -p[1] * u[1]
+    du[2] = -u[1] - p[2] * u[2]
+    return nothing
+end
+
+@testset "Nested ForwardDiff through Rosenbrock Jacobian (#3381)" begin
+    OuterTag = ForwardDiff.Tag{_Nested3381OuterTag, Float64}
+    DualT = ForwardDiff.Dual{OuterTag, Float64, 2}
+    pdual = [
+        DualT(1.5, ForwardDiff.Partials((1.0, 0.0))),
+        DualT(2.0, ForwardDiff.Partials((0.0, 1.0))),
+    ]
+
+    ode_f = ODEFunction{true, SciMLBase.FullSpecialize}(_nested3381_ode!)
+    u0 = [1.0, 1.0]
+    tspan = (0.0, 1.0)
+    prob = ODEProblem(ode_f, u0, tspan, pdual)
+
+    sol = solve(prob, Rosenbrock23(autodiff = AutoForwardDiff(chunksize = 2)))
+    @test SciMLBase.successful_retcode(sol.retcode)
+    @test eltype(sol.u[end]) === DualT
+    # Sanity: primal values still reach the target trajectory.
+    @test all(isfinite, ForwardDiff.value.(sol.u[end]))
+end

lib/OrdinaryDiffEqDifferentiation/test/runtests.jl

@@ -27,6 +27,7 @@ if TEST_GROUP ∉ ("QA", "Sparse", "ModelingToolkit")
     @time @safetestset "Differentiation Trait Tests" include("differentiation_traits_tests.jl")
     @time @safetestset "Autodiff Error Tests" include("autodiff_error_tests.jl")
     @time @safetestset "No Jac Tests" include("nojac_tests.jl")
+    @time @safetestset "Nested ForwardDiff" include("nested_forwarddiff_tests.jl")
 end
 
 # Run sparse tests (separate environment due to ComponentArrays dep conflicts)