Skip to content

Implement forwarddiff_chunksize Val trait for algorithms#3069

Merged
ChrisRackauckas merged 6 commits into
SciML:masterfrom
ChrisRackauckas-Claude:forwarddiff-chunksize-val
Feb 22, 2026
Merged

Implement forwarddiff_chunksize Val trait for algorithms#3069
ChrisRackauckas merged 6 commits into
SciML:masterfrom
ChrisRackauckas-Claude:forwarddiff-chunksize-val

Conversation

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor

Summary

  • Implement SciMLBase.forwarddiff_chunksize trait for all OrdinaryDiffEq algorithm types that use ForwardDiff (implicit, DAE, exponential, composite)
  • Uses existing _get_fwd_chunksize(AD) which returns Val{CS}(), matching the updated trait contract from SciMLBase#1244
  • Single Union dispatch, consistent with existing get_chunksize pattern
  • Bumps SciMLBase compat to 2.143 and DiffEqBase compat to 6.209
  • Fixes @test_broken → plain assignment in dae_rosenbrock_ad_tests.jl since in-place mass matrix solve is now inferable with FunctionWrappers

Dependencies

  • SciMLBase#1244 (merged, v2.143.0 registered)
  • DiffEqBase#1287 (pending CI)

Context

Follow-up to #3066 (merged). The forwarddiff_chunksize trait allows DiffEqBase to compile FunctionWrapper variants with the correct ForwardDiff chunk size, fixing NoFunctionWrapperFoundError on mass matrix ODEProblems when algorithms use non-default chunk sizes.

Test plan

  • CI passes with new SciMLBase and DiffEqBase versions
  • dae_rosenbrock_ad_tests.jl passes (previously @test_broken)

🤖 Generated with Claude Code

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor Author

Added commit cedbd8c which incorporates #3068 with the correct chunk size:

  • _tagged_autodiff(u, ::Val{CS}) uses the algorithm's actual chunk size instead of hardcoding AutoForwardDiff{1}
  • Chunk size from SciMLBase.forwarddiff_chunksize(integrator.alg) threaded through default_nlsolve to all 7 call sites
  • Also includes Pass OrdinaryDiffEqTag to NonlinearSolve in DAE initialization #3068's OrdinaryDiffEqTag fix, OOP argument order fix, and OOP unwrapped_f removal

This PR now supersedes #3068.

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor Author

Fix for CI failures: All 41 failures had the same root cause — AutoForwardDiff{nothing} (the default when no chunk size is specified) was matching _get_fwd_chunksize(::Type{<:AutoForwardDiff{CS}}) where {CS} and returning Val(nothing) instead of Val(0). This caused FunctionWrapper variants to be compiled with Dual{..., nothing} chunk sizes, leading to "No matching function wrapper was found!" errors during precompilation.

Fixed by adding more specific dispatch methods:

_get_fwd_chunksize(::Type{<:AutoForwardDiff{nothing}}) = Val(0)
_get_fwd_chunksize_int(::Type{<:AutoForwardDiff{nothing}}) = 0

Also added _tagged_autodiff(u, ::Val{nothing}) as a safety net in initialize_dae.jl.

Implement the new SciMLBase.forwarddiff_chunksize trait for all
OrdinaryDiffEq algorithm types that use ForwardDiff:

- OrdinaryDiffEqAdaptiveImplicitAlgorithm
- OrdinaryDiffEqImplicitAlgorithm
- DAEAlgorithm
- OrdinaryDiffEqExponentialAlgorithm
- OrdinaryDiffEqAdaptiveExponentialAlgorithm
- CompositeAlgorithm (takes max of sub-algorithms)

This allows DiffEqBase to compile FunctionWrapper variants with the
correct chunk size, fixing NoFunctionWrapperFoundError on mass matrix
ODEProblems when algorithms use non-default chunk sizes.

Also fix @test_broken → plain assignment in dae_rosenbrock_ad_tests.jl
since the in-place mass matrix solve is now inferable with FunctionWrappers.

Requires SciMLBase >= 2.143 and DiffEqBase >= 6.209.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Use _get_fwd_chunksize(AD) which already returns Val{CS}() instead of
_get_fwd_chunksize_int + runtime branching. Consolidates three methods
into one Union dispatch matching existing get_chunksize pattern.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Instead of hardcoding AutoForwardDiff{1}, use the algorithm's actual
ForwardDiff chunk size from SciMLBase.forwarddiff_chunksize. Thread the
chunk size through default_nlsolve and _tagged_autodiff so NonlinearSolve
uses matching Dual types.

Also incorporates PR SciML#3068 fixes:
- Pass OrdinaryDiffEqTag to NonlinearSolve for correct ForwardDiff tag
- Fix 3 OOP default_nlsolve call sites with swapped arguments
- Remove unnecessary unwrapped_f from OOP methods

Closes SciML#3067

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
AutoForwardDiff{nothing} (default when no chunk size is specified) was
matching _get_fwd_chunksize(::Type{<:AutoForwardDiff{CS}}) and returning
Val(nothing) instead of Val(0). This caused FunctionWrapper variants to
be compiled with Dual{..., nothing} chunk sizes, leading to
"No matching function wrapper was found!" errors during precompilation.

Add more specific methods for AutoForwardDiff{nothing} that return
Val(0) / 0 respectively, and add a Val{nothing} handler in
_tagged_autodiff as a safety net.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
The call sites in OrdinaryDiffEqNonlinearSolve now pass 6 args
(including chunksize), but the fallback and error-message overloads
in OrdinaryDiffEqCore only accepted 5 args. This caused a MethodError
when a custom nlsolve algorithm (not Nothing) was passed, since the
6-arg call couldn't match the 5-arg fallback.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
@ChrisRackauckas ChrisRackauckas merged commit 60e8a58 into SciML:master Feb 22, 2026
260 of 275 checks passed
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