MTKBaseEnzymeExt: mirror Mooncake's tangent_type=NoTangent via EnzymeRules.inactive_type#4558
Closed
ChrisRackauckas-Claude wants to merge 1 commit into
Conversation
…rations
Adds a new package extension `MTKBaseEnzymeExt` that registers
`EnzymeRules.inactive_type` for the same MTKBase-owned types that
`MTKMooncakeExt` marks `Mooncake.tangent_type = NoTangent`:
- `System` (also covered by the existing `<:AbstractSystem` declaration
in `ModelingToolkitBase.jl`; kept here for one-to-one parity with the
Mooncake extension and for explicitness),
- `GeneratedFunctionWrapper` — the `RuntimeGeneratedFunction` wrapper
used as the model RHS,
- `ObservedFunctionCache` — holds a `System` reference and a
`Dict{Any,Any}` of cached observed-function closures,
- `MissingGuessValue.Type` — Moshi `@data` tagged union used as a
configuration enum for the initialization solver.
These types are non-differentiable structural / configuration data.
Without inactive_type declarations, Enzyme's runtime-activity dispatch
and `make_zero` traversal recurse into self-referential or
otherwise-unhandleable fields on closures / problem / solution structs
that transitively capture them, producing activity-analysis errors or
`MethodError`s in `create_activity_wrapper`.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Author
|
Closing as not required. Ablation study with the desauty SCC init test rewritten to use the proper Enzyme API ( |
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
Adds a new package extension
MTKBaseEnzymeExt(triggered onEnzyme) that registersEnzymeRules.inactive_typefor the same MTKBase-owned types thatMTKMooncakeExtalready marksMooncake.tangent_type = NoTangent. This is one of several follow-ups for full Enzyme support through MTK-generated problem / initialization-problem closures.Types declared inactive
Systemlib/ModelingToolkitBase/src/systems/system.jlsystems::Vector{System},parent::Union{Nothing, System},initializesystem::Union{Nothing, System}. Already covered by the existing<:AbstractSystemdeclaration inModelingToolkitBase.jl; restated here for one-to-one parity withMTKMooncakeExtand explicitness.GeneratedFunctionWrapperlib/ModelingToolkitBase/src/systems/codegen_utils.jlRuntimeGeneratedFunctions— the model RHS closure, not a derivative carrier.ObservedFunctionCachelib/ModelingToolkitBase/src/systems/abstractsystem.jlSystemreference and aDict{Any, Any}of cached observed-function closures.MissingGuessValue.Typelib/ModelingToolkitBase/src/systems/problem_utils.jl@datatagged union (Constant{<:Number} | Random{<:AbstractRNG} | Error) used as a configuration enum for the initialization solver.Why this is needed
Without these declarations, Enzyme's runtime-activity dispatch and
make_zerotraversal recurse into self-referential or otherwise-unhandleable fields when a closure (or problem / solution struct) transitively captures one of these values, producing activity-analysis errors orMethodErrors increate_activity_wrapper.MTKMooncakeExtalready handles the equivalent on the Mooncake side viaMooncake.tangent_type(::Type{T}) = NoTangent. This PR brings the Enzyme side to parity for the types MTKBase owns. Types that live in other packages (SciMLBase'sOverrideInitData/ODENLStepData,FunctionWrappersWrappers'FunctionWrappersWrapper,Base'sImmutableDict) are intentionally not declared here — they belong in their own packages' Enzyme extensions; separate PRs to follow.Verification
Local smoke test confirmed the extension loads and the declarations apply:
Project.toml changes
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"to[weakdeps].MTKBaseEnzymeExt = "Enzyme"to[extensions].Enzyme = "0.13"to[compat].ModelingToolkitBasepatch version to1.40.1.Notes
Please ignore until reviewed by @ChrisRackauckas.
This is the second follow-up in a series of small PRs targeting MTKBase's Enzyme story. The first was the
<:AbstractSysteminactive_typedeclaration already on master. A separate PR will address theMTKParameters.canonicalize(::Tunable, ...)repack closure aliasing freshcachesbuffers — that fix is independent and stacked on its own branch (enzyme-fresh-caches-on-repack).Test plan
test_caches_fix.jlMWE referenced in the surrounding investigation).Co-Authored-By: Chris Rackauckas accounts@chrisrackauckas.com