DiffEqBase: skip AutoSpecialize FWW wrap when u0 is Dual#3642
Closed
ChrisRackauckas-Claude wants to merge 1 commit into
Closed
DiffEqBase: skip AutoSpecialize FWW wrap when u0 is Dual#3642ChrisRackauckas-Claude wants to merge 1 commit into
ChrisRackauckas-Claude wants to merge 1 commit into
Conversation
The AutoSpecialize path in `promote_f` wraps the user RHS in a
`FunctionWrappersWrapper` keyed off `(u0, u0, p, t)` types so that
similar-typed problems share compiled specializations. The FWW signature
slots are baked at construction time from the seen `u0` type.
When `u0` is already a `Vector{<:ForwardDiff.Dual}` (because the solve
is happening inside an outer ForwardDiff layer — e.g. `gradient` /
`hessian` over a `solve(NLLS)` or `solve(ODE)` call), the precompiled
slot doesn't match what the solver actually evaluates: solvers like
Rosenbrock re-widen `u`/`p` into deeper-nested Dual types for the inner
Jacobian, and there is no FWW slot for that deeper type, producing
`NoFunctionWrapperFoundError`, or worse, silent dispatch into a slot
whose compiled body builds Duals at the wrong nesting depth.
Skip the wrap entirely when `isdualtype(eltype(u0))`. With no wrapper,
Julia compiles a fresh specialization of the user RHS on the actual
nested-Dual types each time, which is the same code path
`FullSpecialize` would take. There is no cost: the type-erasure benefit
of AutoSpecialize only matters across many similarly-shaped non-Dual
problems, which is not the use case when `u0` is a one-off Dual array
produced by an outer differentiation pass.
Bumps DiffEqBase 7.5.0 -> 7.5.1.
Adds a regression test in `test/downstream/unwrapping.jl` covering both
directions:
- `Dual` u0 + AutoSpecialize: `integ.f.f === g!` (no FWW wrap).
- `Float64` u0 + AutoSpecialize: `integ.f.f isa FunctionWrappersWrapper`
(wrap path preserved).
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Contributor
Author
|
Closing — empirical testing shows this patch alone does not fix the #3381 MWE. The crash originates in the cross-tag multiply inside the Rosenbrock inner Jacobian (not in the FWW wrapping path the patch targets), so skipping the AutoSpecialize wrap for Dual u0 has no observable effect: stock vs patched DiffEqBase both fail on stock ForwardDiff and both succeed under the hybrid |
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
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.
Summary
Adds a single guard to
promote_finDiffEqBase/src/solve.jl: when the AutoSpecialize path would normally wrap the user RHS in aFunctionWrappersWrapper, skip the wrap ifeltype(u0) <: ForwardDiff.Dual.Why
The FWW signature slots are baked from
(u0, u0, p, t)types at construction time. Whenu0is alreadyVector{<:Dual}(the solve is happening inside an outer ForwardDiff layer —gradient/hessianover asolve(...)call), the precompiled slot doesn't match what the solver actually evaluates: solvers like Rosenbrock re-widenu/pinto deeper-nested Dual types for the inner Jacobian, and there is no FWW slot for that deeper type. The result is eitherNoFunctionWrapperFoundErroror silent dispatch into a slot whose compiled body builds Duals at the wrong nesting depth (the latter is part of the failure shape behind #3381 / #3587's AutoSpecialize path).With no wrapper, Julia compiles a fresh specialization on the actual nested-Dual types each call, same as
FullSpecialize. The type-erasure benefit of AutoSpecialize only matters across many similarly-shaped non-Dual problems; that's not the use case whenu0is a one-off Dual array produced by an outer differentiation pass.Diff
lib/DiffEqBase/src/solve.jl: add!SciMLBase.isdualtype(eltype(u0))to the AutoSpecialize branch in bothpromote_fpaths (uses-ForwardDiff and no-ForwardDiff).lib/DiffEqBase/Project.toml: bump 7.5.0 → 7.5.1.lib/DiffEqBase/test/downstream/unwrapping.jl: regression test pair —Dualu0 + AutoSpecialize:integ.f.f === g!(no FWW wrap).Float64u0 + AutoSpecialize:integ.f.f isa FunctionWrappersWrapper(wrap path preserved).Test plan
test/downstream/unwrapping.jlpass (Julia 1.11, Downstream env).integ.f.f === g!for Dual u0 andinteg.f.f isa FWWfor Float64 u0 withAutoSpecialize.Note on scope
This is independent of #3587. #3587 fixes the cross-tag-multiply nesting issue inside Rosenbrock's inner Jacobian (via the
_widen_uf_p_for_jacstep). This PR fixes a different failure mode — FWW slot mismatch whenu0is already a Dual. The two are complementary: even with #3587 merged, the FWW slot issue can still bite users who constructODEFunction{true, AutoSpecialize}(...)and then solve under an outer ForwardDiff layer.🤖 Generated with Claude Code