Fix inference issue in Rosenbrock DAE test #2931
Merged
+132
−68
Conversation
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
The alg_autodiff function was type-unstable because it used runtime
conditionals (if/elseif/else) that could return different types:
- AutoForwardDiff() for Val{true}
- AutoFiniteDiff() for Val{false}
- The original autodiff value otherwise
This caused downstream type inference failures in build_jac_config and
build_grad_config, which in turn caused the cache type and solution type
to be uninferred in the integrator construction.
The fix uses dispatch-based helper functions (_unwrap_autodiff) instead
of runtime conditionals, which allows Julia to resolve the return type
at compile time based on the input type.
This is the first step toward fixing inference for Rosenbrock DAE tests.
Comment on lines
-22
to
-28
| if autodiff == Val(true) | ||
| return AutoForwardDiff() | ||
| elseif autodiff == Val(false) | ||
| return AutoFiniteDiff() | ||
| else | ||
| return autodiff | ||
| end |
Member
There was a problem hiding this comment.
alternatively, this could be
if autodiff isa Val{true}
return AutoForwardDiff()
elseif autodiff isa Val{false}
return AutoFiniteDiff()
else
return autodiff
end
Contributor
There was a problem hiding this comment.
Is there any true advantage of using branching over multiple dispatch for distinguishing executable code by types? I know this is very widely used in the SciML ecosystem, but this if-else by type design pattern causes a lot of trouble when trying to extend codes.
Member
Author
There was a problem hiding this comment.
we can do splits to multiple dispatch, it's just a code structure thing.
Use dispatch-based patterns instead of runtime conditionals to enable
better type inference. Both functions now dispatch on Val{true/false}
based on whether config computation is needed, which allows Julia's
compiler to specialize each branch separately.
This helps the cache construction in Rosenbrock methods infer concrete
types for jac_config and grad_config fields.
Refactor build_jac_config and build_grad_config to dispatch on the actual field values (f.jac, f.Wfact_t, f.tgrad) rather than using Val(boolean_condition). This allows Julia to specialize based on whether these fields have type Nothing vs a function type. Also add function barrier _make_rosenbrock_cache to help ensure cache construction sees concrete types.
Contributor
Benchmark Results (Julia v1)Time benchmarks
Memory benchmarks
|
Convert remaining runtime conditional to type-based dispatch: - _build_jac_config_alg(::Nothing, ::Nothing) -> compute config - _build_jac_config_alg(::Nothing, linsolve) -> check needs_concrete_A - _build_jac_config_alg(::Bool, linsolve) -> check bool value For default Rodas5P (concrete_jac=nothing, linsolve=nothing), this dispatches directly to the config computation without any runtime conditionals, enabling full type inference.
ChrisRackauckas
deleted the
claude/fix-rosenbrock-inference-01HJCdwyyMxVHCrV6qJHPqGj
branch
Member
Author
|
accidentally merged and reverted |
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.
The alg_autodiff function was type-unstable because it used runtime conditionals (if/elseif/else) that could return different types:
This caused downstream type inference failures in build_jac_config and build_grad_config, which in turn caused the cache type and solution type to be uninferred in the integrator construction.
The fix uses dispatch-based helper functions (_unwrap_autodiff) instead of runtime conditionals, which allows Julia to resolve the return type at compile time based on the input type.
This is the first step toward fixing inference for Rosenbrock DAE tests.