Skip to content

Add MixedDuplicated(::T, ::T) shim for Enzyme#3126 workaround#1459

Draft
ChrisRackauckas-Claude wants to merge 2 commits into
SciML:masterfrom
ChrisRackauckas-Claude:enzyme-mixed-shim
Draft

Add MixedDuplicated(::T, ::T) shim for Enzyme#3126 workaround#1459
ChrisRackauckas-Claude wants to merge 2 commits into
SciML:masterfrom
ChrisRackauckas-Claude:enzyme-mixed-shim

Conversation

@ChrisRackauckas-Claude
Copy link
Copy Markdown
Contributor

Please ignore until reviewed by @ChrisRackauckas.

Summary

Enzyme's runtime_generic_augfwdcreate_activity_wrapper (Enzyme/src/rules/jitrules.jl line 14) constructs MixedDuplicated(primarg, shadowarg) with shadowarg::T rather than Base.RefValue{T}. EnzymeCore.MixedDuplicated only defines MixedDuplicated(::T, ::Base.RefValue{T}), so the bare-T call raises a MethodError. This blocks reverse-mode AD chains that capture ODESolution (e.g. the MTK parameter-init path) — exactly the residual blocker tracked in EnzymeAD/Enzyme.jl#3126.

Until #3126 lands upstream, this PR adds a workaround:

  • A type-piracy shim in src/enzyme_rules.jl that intercepts MixedDuplicated(x::T, dx::T, check::Bool=true) and delegates to the existing constructor by wrapping the shadow in Base.RefValue. The constructor's check::Bool flag is currently unused by EnzymeCore, so the shim is invariant-preserving.
  • Flips two @test_broken blocks blocked specifically on this MethodError back to @test:
    • test/parameter_initialization.jl: "Adjoint through Prob (Enzyme)"
    • test/mtk.jl: MTK Forward Mode Enzyme block (line 204)
  • Mooncake @test_broken in test/mtk.jl (line 230) is unrelated and left alone.

Risk acknowledgment (type piracy)

This is type piracy on EnzymeCore. Concerns:

  • The shim adds a method to EnzymeCore.MixedDuplicated that fires for any (::T, ::T) call site. If Enzyme (or downstream) ever calls MixedDuplicated(::T, ::T) with a legitimate intent (e.g. a T = Base.RefValue{S} primal whose shadow really is a bare RefValue{S}), this shim silently double-wraps and the result is wrong. I have not found such a call site in the current Enzyme tree — every other MixedDuplicated construction either explicitly wraps in Base.RefValue or goes through the broken create_activity_wrapper path — but the surface is wide.
  • The remedy if it bites: remove this shim and the two @test_broken flips, return to @test_broken, and wait for Enzyme#3126 to land upstream.

CI will surface any breakage of other Enzyme code paths via the full test matrix.

Verification

I attempted to run Pkg.test() on Julia 1.12 with GROUP=Core8 locally. SciMLSensitivity precompiles cleanly with the shim. The actual parameter_initialization.jl test entered the Enzyme reverse-mode compilation phase and was still grinding through code generation past 25 minutes elapsed (with 4+ GB RSS) when I had to push for CI verification. I did NOT see a final pass/fail locally. CI here is the authoritative verification — please review CI results before merging.

Removal plan

When EnzymeAD/Enzyme.jl#3126 lands upstream, delete:

  • The MixedDuplicated(x::T, dx::T, ...) method in src/enzyme_rules.jl
  • The import Enzyme.EnzymeCore: MixedDuplicated line

Tests should continue to pass on the fixed Enzyme version.

🤖 Generated with Claude Code

ChrisRackauckas and others added 2 commits May 25, 2026 23:00
…vity

The `EnzymeOriginator` overload at concrete_solve.jl:377 called
`Enzyme.gradient(Enzyme.Reverse, Const(f), tunables)` without
`set_runtime_activity`. The outer user-level call's
`set_runtime_activity(Reverse)` flag does not propagate into this
nested gradient, so the inner Enzyme call raised
`EnzymeRuntimeActivityError` at MTK's init `remake` even when the
caller used the documented Enzyme pattern.

Wrap with `set_runtime_activity(Reverse)` so the nested gradient
respects the same activity relaxation as the outer call. Unblocks the
`@test_broken` Enzyme block in `test/parameter_initialization.jl`
(SciML#1455).

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Enzyme's `runtime_generic_augfwd` (`create_activity_wrapper` in
`jitrules.jl`) constructs `MixedDuplicated(primarg, shadowarg)` with the
shadow as a bare `T` rather than `Base.RefValue{T}`. The only constructor
defined in `EnzymeCore.MixedDuplicated` requires `Base.RefValue{T}` for
the shadow, so the bare-`T` call raises a MethodError, breaking
reverse-mode AD chains that capture `ODESolution` (e.g. the MTK
parameter-init path).

Add a type-piracy shim in `src/enzyme_rules.jl` that wraps the shadow in
`Base.RefValue` before delegating to the upstream constructor. The
existing constructor's `check::Bool` parameter is forwarded but is
currently unused by EnzymeCore, so the shim is invariant-preserving.

Flip the two `@test_broken` blocks that were blocked specifically on this
MethodError back to `@test`:
- `test/parameter_initialization.jl` (Adjoint through Prob via Enzyme)
- `test/mtk.jl` (MTK Forward Mode Enzyme block, line 204)

Remove the shim once EnzymeAD/Enzyme.jl#3126 lands upstream.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Copy link
Copy Markdown
Contributor

@wsmoses wsmoses left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is absolutely wrong and type piracy and should not be done

@ChrisRackauckas
Copy link
Copy Markdown
Member

I know it's just testing CI to see if it would make the test pass.

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.

3 participants