Add Enzyme inactive rules for VJP choice types #1272
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
Adds proper Enzyme extension for SciMLSensitivity to make VJP choice types inactive during differentiation. This complements the AbstractSensitivityAlgorithm rule in SciMLBase. Fixes SciMLSensitivity.jl#1225: When `sensealg=BacksolveAdjoint(autojacvec=EnzymeVJP())` is passed to ODEProblem constructor, Enzyme would fail with "Error handling recursive stores for String" because it tried to differentiate through the VJP choice objects. Changes: - Add ext/SciMLSensitivityEnzymeExt.jl with VJPChoice inactive rule - Add comprehensive tests in test/enzyme_vjp_inactive.jl - Tests include original failing case from issue SciML#1225 This works in conjunction with SciMLBase PR that adds AbstractSensitivityAlgorithm inactive rule to avoid type piracy. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
This version includes the Enzyme extension for AbstractSensitivityAlgorithm inactive_type rule, which is needed for this PR's VJP choice rules to work together to fix SciMLSensitivity.jl#1225. Corresponds to General registry PR #138113. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Run JuliaFormatter with SciMLStyle on the new files to fix format check errors. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Since SciMLSensitivity directly depends on Enzyme, use direct include instead of extension for VJP choice inactive rules. Changes: - Remove ext/SciMLSensitivityEnzymeExt.jl - Add src/enzyme_rules.jl with VJPChoice inactive rule - Include in main SciMLSensitivity.jl module - Remove extension registration from Project.toml Note: This addresses the original "Error handling recursive stores for String" error from SciML#1225, but there may be additional Enzyme compatibility issues with array operations as noted in recent issue comments. 🤖 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
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 Enzyme extension to make VJP choice types inactive during differentiation, complementing the AbstractSensitivityAlgorithm rule in SciMLBase.
Problem
When
sensealg=BacksolveAdjoint(autojacvec=EnzymeVJP())is passed toODEProblemconstructor, Enzyme fails with:This occurs because Enzyme tries to differentiate through VJP choice objects (like
EnzymeVJP(),ZygoteVJP(), etc.) that are stored within the sensitivity algorithm.See issue: #1225
Solution
This PR adds:
ext/SciMLSensitivityEnzymeExt.jlwith inactive rule for VJP choice typestest/enzyme_vjp_inactive.jlincluding the original failing case from Enzyme gradient fails whensensealgis defined in theODEProblemrather than in thesolvefunction #1225Architecture
This PR handles VJP choice types defined in SciMLSensitivity, while the
AbstractSensitivityAlgorithminactive rule is properly handled in SciMLBase to avoid type piracy.Complete solution requires both:
AbstractSensitivityAlgorithminactive ruleVJPChoiceinactive rule (this PR)Testing
EnzymeVJP,ZygoteVJP,ReverseDiffVJP,TrackerVJP)sensealgis defined in theODEProblemrather than in thesolvefunction #1225 withsensealginODEProblemconstructorBefore/After
Before:
After:
Both cases work correctly with Enzyme differentiation.
🤖 Generated with Claude Code