Skip to content

MTKBaseEnzymeExt: mirror Mooncake's tangent_type=NoTangent via EnzymeRules.inactive_type#4558

Closed
ChrisRackauckas-Claude wants to merge 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:mtkbase-enzyme-inactive-type-mirror-v2
Closed

MTKBaseEnzymeExt: mirror Mooncake's tangent_type=NoTangent via EnzymeRules.inactive_type#4558
ChrisRackauckas-Claude wants to merge 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:mtkbase-enzyme-inactive-type-mirror-v2

Conversation

@ChrisRackauckas-Claude
Copy link
Copy Markdown

Summary

Adds a new package extension MTKBaseEnzymeExt (triggered on Enzyme) that registers EnzymeRules.inactive_type for the same MTKBase-owned types that MTKMooncakeExt already marks Mooncake.tangent_type = NoTangent. This is one of several follow-ups for full Enzyme support through MTK-generated problem / initialization-problem closures.

Types declared inactive

Type Where it lives Rationale
System lib/ModelingToolkitBase/src/systems/system.jl Self-referential fields systems::Vector{System}, parent::Union{Nothing, System}, initializesystem::Union{Nothing, System}. Already covered by the existing <:AbstractSystem declaration in ModelingToolkitBase.jl; restated here for one-to-one parity with MTKMooncakeExt and explicitness.
GeneratedFunctionWrapper lib/ModelingToolkitBase/src/systems/codegen_utils.jl Wraps RuntimeGeneratedFunctions — the model RHS closure, not a derivative carrier.
ObservedFunctionCache lib/ModelingToolkitBase/src/systems/abstractsystem.jl Contains a System reference and a Dict{Any, Any} of cached observed-function closures.
MissingGuessValue.Type lib/ModelingToolkitBase/src/systems/problem_utils.jl Moshi @data tagged 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_zero traversal 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 or MethodErrors in create_activity_wrapper.

MTKMooncakeExt already handles the equivalent on the Mooncake side via Mooncake.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's OverrideInitData / ODENLStepData, FunctionWrappersWrappers' FunctionWrappersWrapper, Base's ImmutableDict) 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:

System inactive: true
GFW{(3,3,true),Nothing,Nothing} inactive: true
ObservedFunctionCache inactive: true
MissingGuessValue.Type inactive: true
MissingGuessValue.Error type inactive: true

Project.toml changes

  • Adds Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" to [weakdeps].
  • Adds MTKBaseEnzymeExt = "Enzyme" to [extensions].
  • Adds Enzyme = "0.13" to [compat].
  • Bumps ModelingToolkitBase patch version to 1.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 <:AbstractSystem inactive_type declaration already on master. A separate PR will address the MTKParameters.canonicalize(::Tunable, ...) repack closure aliasing fresh caches buffers — that fix is independent and stacked on its own branch (enzyme-fresh-caches-on-repack).

Test plan

  • CI green on MTKBase tests (no regression from the inactive_type declarations).
  • Downstream smoke: SCC init-problem gradient through Enzyme (see test_caches_fix.jl MWE referenced in the surrounding investigation).

Co-Authored-By: Chris Rackauckas accounts@chrisrackauckas.com

…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>
@ChrisRackauckas-Claude
Copy link
Copy Markdown
Author

Closing as not required. Ablation study with the desauty SCC init test rewritten to use the proper Enzyme API (Enzyme.autodiff + Duplicated(iprob), SciML/SciMLSensitivity.jl#1454) confirms this PR doesn't change the outcome — desauty still passes with MTKBaseEnzymeExt disabled (file removed and [extensions] entry commented out). The Mooncake-parity inactive_type declarations may still be useful for other Enzyme paths, but I don't have a concrete test that needs them. Can be reopened if one comes up.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants