DiffEqBaseEnzymeExt: tighten sensealg inner type to exclude SensitivityADPassThrough (#3681)

PR #3678 loosened the `augmented_primal`/`reverse` rules from
`Const{Nothing}|Const{<:AbstractSensitivityAlgorithm}` to a plain
`Enzyme.Annotation` and added a runtime check throwing on anything
other than `nothing`/`AbstractSensitivityAlgorithm`. That broke the
existing `Direct Differentiation of Explicit ODE Solve` AD test
(test/ad/ad_tests.jl:583), which passes
`sensealg = SensitivityADPassThrough()` to deliberately bypass adjoint
machinery and do direct AD through `solve_up`. Pre-#3678 the rule
simply did not dispatch for `Const{SensitivityADPassThrough}` (since
`SensitivityADPassThrough <: AbstractDEAlgorithm`, not
`<: AbstractSensitivityAlgorithm`), so Enzyme fell through to direct
AD. Post-#3678 the rule swallowed it and immediately raised.

Restore the original dispatch behavior while keeping #3678's wider
outer-annotation match: constrain the *inner* sensealg value type to
`Union{Nothing, AbstractSensitivityAlgorithm}` and drop the redundant
runtime ArgumentError check. The outer `Enzyme.Annotation` (Const,
Duplicated, MixedDuplicated, Active) is still freely accepted, which
preserves the original motivation of #3678 (tolerate
`set_runtime_activity(Reverse)` promoting a sensealg const to
`Duplicated{GaussAdjoint{...}}` before the rule matches).

Bumps DiffEqBase to 7.5.3.

Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>

Author: ChrisRackauckas-Claude

lib/DiffEqBase/Project.toml

@@ -1,7 +1,7 @@
 name = "DiffEqBase"
 uuid = "2b5f629d-d688-5b77-993f-72d75c75574e"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "7.5.2"
+version = "7.5.3"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

lib/DiffEqBase/ext/DiffEqBaseEnzymeExt.jl

@@ -34,19 +34,18 @@ module DiffEqBaseEnzymeExt
     # can still promote module-level `const _MTK_SENSEALG = GaussAdjoint(...)`
     # references — or any sensealg captured in an `ODEProblem` field — to
     # `Duplicated{<:AbstractSensitivityAlgorithm}` before this rule is matched.
-    # Accept any `Enzyme.Annotation` (Const / Duplicated / MixedDuplicated /
-    # Active) here and read only `sensealg.val`; the shadow has no meaning.
+    # Accept any `Enzyme.Annotation` wrapping a sensealg value (Const /
+    # Duplicated / MixedDuplicated / Active) and read only `sensealg.val`;
+    # the shadow has no meaning. The inner-type constraint keeps
+    # `SensitivityADPassThrough` (which is `<:AbstractDEAlgorithm`, not
+    # `<:AbstractSensitivityAlgorithm`) off this rule so that case still
+    # falls through to direct AD of `solve_up`.
     function Enzyme.EnzymeRules.augmented_primal(
             config::Enzyme.EnzymeRules.RevConfigWidth{1},
             func::Const{typeof(DiffEqBase.solve_up)}, RTA::Union{Type{Duplicated{RT}}, Type{MixedDuplicated{RT}}}, prob,
-            sensealg::Enzyme.Annotation,
+            sensealg::Enzyme.Annotation{<:Union{Nothing, DiffEqBase.AbstractSensitivityAlgorithm}},
             u0, p, args...; kwargs...
         ) where {RT}
-        sensealg.val isa Union{Nothing, DiffEqBase.AbstractSensitivityAlgorithm} ||
-            throw(ArgumentError(
-            "DiffEqBaseEnzymeExt: `sensealg` must be `nothing` or " *
-                "`<:AbstractSensitivityAlgorithm`, got $(typeof(sensealg.val))",
-        ))
 
         res = DiffEqBase._solve_adjoint(
             copy_or_reuse(config, prob.val, 2), copy_or_reuse(config, sensealg.val, 3),
@@ -77,7 +76,7 @@ module DiffEqBaseEnzymeExt
     function Enzyme.EnzymeRules.reverse(
             config::Enzyme.EnzymeRules.RevConfigWidth{1},
             func::Const{typeof(DiffEqBase.solve_up)}, ::Union{Type{Duplicated{RT}}, Type{MixedDuplicated{RT}}}, tape, prob,
-            sensealg::Enzyme.Annotation,
+            sensealg::Enzyme.Annotation{<:Union{Nothing, DiffEqBase.AbstractSensitivityAlgorithm}},
             u0, p, args...; kwargs...
         ) where {RT}