Implement forwarddiff_chunksize Val trait for algorithms#3069
Merged
ChrisRackauckas merged 6 commits intoFeb 22, 2026
Merged
Conversation
Contributor
Author
|
Added commit
This PR now supersedes #3068. |
2 tasks
Merged
Contributor
Author
|
Fix for CI failures: All 41 failures had the same root cause — Fixed by adding more specific dispatch methods: _get_fwd_chunksize(::Type{<:AutoForwardDiff{nothing}}) = Val(0)
_get_fwd_chunksize_int(::Type{<:AutoForwardDiff{nothing}}) = 0Also added |
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>
1c739c0 to
36c2dc8
Compare
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>
5 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
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
SciMLBase.forwarddiff_chunksizetrait for all OrdinaryDiffEq algorithm types that use ForwardDiff (implicit, DAE, exponential, composite)_get_fwd_chunksize(AD)which returnsVal{CS}(), matching the updated trait contract from SciMLBase#1244get_chunksizepattern@test_broken→ plain assignment indae_rosenbrock_ad_tests.jlsince in-place mass matrix solve is now inferable with FunctionWrappersDependencies
Context
Follow-up to #3066 (merged). The
forwarddiff_chunksizetrait allows DiffEqBase to compile FunctionWrapper variants with the correct ForwardDiff chunk size, fixingNoFunctionWrapperFoundErroron mass matrix ODEProblems when algorithms use non-default chunk sizes.Test plan
dae_rosenbrock_ad_tests.jlpasses (previously@test_broken)🤖 Generated with Claude Code