Fix nested ForwardDiff tag through wrapfun_iip + UJacobianWrapper#3412
Closed
ChrisRackauckas-Claude wants to merge 1 commit intoSciML:masterfrom
Closed
Conversation
…ciML#3381) When a Rosenbrock integrator runs with a `Vector{<:Dual}` `p` (i.e. 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 prepared `JacobianConfig`, but `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 and dispatches through ForwardDiff's `@generated tagcount` precedence — whose literal value is baked at first compile and varies with precompile ordering. That ordering can invert the nesting and crash the eventual `setindex!(du, ...)` with `MethodError: no method matching Float64(::nested_dual)`. Fix at the inner solver's side by lifting `uf.p` into the inner nested-Dual type once, before delegating to `DI.jacobian!`. The widened `p` carries zero inner partials (`p` does not depend on `u`). One `convert.(inner_T, p)` allocation per `jacobian!` call is amortized across every chunk evaluation DI performs in that call — no per-call allocation inside the hot loop, unlike a `UJacobianWrapper`-level override. Skips widening when `uf.f` holds a `FunctionWrappersWrapper` (`AutoSpecialize`): DiffEqBase already precompiles the `(nested_u, outer_p, t)` FW signature directly, and a widened `p` would produce `(nested_u, nested_p, t)` that matches no precompiled slot. Adds `test/nested_forwarddiff_tests.jl`, exercising both the MRE call graph and a hand-rolled outer tag (the latter deterministically reproduces the crash pre-fix). Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
ChrisRackauckas-Claude
force-pushed
the
fix-nested-forwarddiff-tag-in-wrapfun
branch
from
0e5b4f9 to
1773ee0
Compare
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fix for #3381. Alternative to #3406 addressing the review feedback at #3406 (comment) — handles the nested tag at the source (
wrapfun_iip+UJacobianWrapper) instead of unwrapping the FunctionWrappersWrapper.Summary
wrapfun_iipnow usesSciMLBase.anyeltypedualto detect a Dual-carryingp, promotes the innerTsodualgen(T)yields the properly nested dual, and promotes the compiled signature'sptype for Jacobian variants.UJacobianWrapper{true}(du1, uprev::AbstractArray{<:Dual{<:Any,<:Dual}})override eagerly liftspinto the seeded DualT sou*pstays within one tag hierarchy and FunctionWrappersWrapper dispatches on the compiled nested-Dual signature (no raw-call fallback).Root cause
NonlinearSolveseedspwithDual{NonlinearSolveTag, Float64, 2}, then insideresid!the ODE solve seedsuwithDual{OrdinaryDiffEqTag, Dual{NonlinearSolveTag, Float64, 2}, CS}. Before this patchwrapfun_iipcompiled the Jacobian signature withp::Vector{Dual{NLTag}}(the original T3), so in the user functionu*pcrossed sibling tag hierarchies — ForwardDiff's tag precedence returned a triple-nestedDual{NLTag, Dual{OrdEqTag, …}, …}that cannot be stored back intodu::Vector{Dual{OrdEqTag, …, CS}}, producing:Test plan
ODEDIFFEQ_TEST_GROUP=Core Pkg.test(\"DiffEqBase\")passes locallyMethodErrorabove when the patch is removeddual_detection_solution.jlregression test🤖 Generated with Claude Code