Add Enzyme extension to make AbstractSensitivityAlgorithm inactive #1137
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
Fixes SciMLSensitivity.jl#1225: When `sensealg` is passed to ODEProblem constructor, Enzyme would fail trying to differentiate through the sensitivity algorithm object. This adds an inactive_type rule for AbstractSensitivityAlgorithm to treat it as constant during Enzyme differentiation. This is the correct place for this fix since AbstractSensitivityAlgorithm is defined in SciMLBase, avoiding type piracy that would occur if the rule were defined in SciMLSensitivity. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
EnzymeRules is the correct package to import for defining inactive_type rules. Updated Project.toml to use EnzymeRules as weakdep instead of Enzyme. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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 an Enzyme extension to SciMLBase to make
AbstractSensitivityAlgorithminactive during Enzyme differentiation.Problem
Currently, when sensitivity algorithms are passed to
ODEProblemconstructors (e.g.,ODEProblem(f, u0, tspan, p; sensealg=BacksolveAdjoint(autojacvec=EnzymeVJP()))), Enzyme fails with:This occurs because Enzyme tries to differentiate through the sensitivity algorithm object, which contains complex structures that Enzyme cannot handle.
See issue: SciML/SciMLSensitivity.jl#1225
Solution
This PR adds a proper Enzyme extension to SciMLBase with an
inactive_typerule:This is the correct location for this fix because:
AbstractSensitivityAlgorithmis defined in SciMLBaseChanges
ext/SciMLBaseEnzymeExt.jlwith the inactive type ruleTesting
The fix resolves the original issue where:
solve(prob, alg; sensealg=BacksolveAdjoint(...))worksODEProblem(f, u0, tspan, p; sensealg=BacksolveAdjoint(...))failsAfter this change, both work correctly with Enzyme differentiation.
🤖 Generated with Claude Code