Skip to content

DiffEqBase: skip AutoSpecialize FWW wrap when u0 is Dual#3642

Closed
ChrisRackauckas-Claude wants to merge 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:fix/diffeqbase-skip-autospecialize-dual-u0
Closed

DiffEqBase: skip AutoSpecialize FWW wrap when u0 is Dual#3642
ChrisRackauckas-Claude wants to merge 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:fix/diffeqbase-skip-autospecialize-dual-u0

Conversation

@ChrisRackauckas-Claude
Copy link
Copy Markdown
Contributor

Draft — please ignore until reviewed by @ChrisRackauckas.

Summary

Adds a single guard to promote_f in DiffEqBase/src/solve.jl: when the AutoSpecialize path would normally wrap the user RHS in a FunctionWrappersWrapper, skip the wrap if eltype(u0) <: ForwardDiff.Dual.

Why

The FWW signature slots are baked from (u0, u0, p, t) types at construction time. When u0 is already Vector{<:Dual} (the solve is happening inside an outer ForwardDiff layer — gradient/hessian over a solve(...) 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. The result is either NoFunctionWrapperFoundError or 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 when u0 is 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 both promote_f paths (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 —
    • Dual u0 + AutoSpecialize: integ.f.f === g! (no FWW wrap).
    • Float64 u0 + AutoSpecialize: integ.f.f isa FunctionWrappersWrapper (wrap path preserved).

Test plan

  • Local: new tests in test/downstream/unwrapping.jl pass (Julia 1.11, Downstream env).
  • Manual: confirmed integ.f.f === g! for Dual u0 and integ.f.f isa FWW for Float64 u0 with AutoSpecialize.
  • CI green.

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_jac step). This PR fixes a different failure mode — FWW slot mismatch when u0 is already a Dual. The two are complementary: even with #3587 merged, the FWW slot issue can still bite users who construct ODEFunction{true, AutoSpecialize}(...) and then solve under an outer ForwardDiff layer.

🤖 Generated with Claude Code

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>
@ChrisRackauckas-Claude ChrisRackauckas-Claude deleted the fix/diffeqbase-skip-autospecialize-dual-u0 branch May 17, 2026 13:27
@ChrisRackauckas-Claude
Copy link
Copy Markdown
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 fix tested separately. Not worth merging as a no-op.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants