WIP add EnzymeAdjoint

Fixes #1148. Needs tests, and it needs the actual direct adjoints of Enzyme to work again

Author: ChrisRackauckas

src/SciMLSensitivity.jl

@@ -95,7 +95,7 @@ export ODEForwardSensitivityFunction, ODEForwardSensitivityProblem, SensitivityF
        shadow_forward, shadow_adjoint
 
 export BacksolveAdjoint, QuadratureAdjoint, GaussAdjoint, InterpolatingAdjoint,
-       TrackerAdjoint, ZygoteAdjoint, ReverseDiffAdjoint,
+       TrackerAdjoint, ZygoteAdjoint, EnzymeAdjoint, ReverseDiffAdjoint,
        ForwardSensitivity, ForwardDiffSensitivity,
        ForwardDiffOverAdjoint,
        SteadyStateAdjoint,

src/concrete_solve.jl

@@ -1258,6 +1258,88 @@ function DiffEqBase._concrete_solve_adjoint(
         p)
 end
 
+function DiffEqBase._concrete_solve_adjoint(
+        prob::Union{SciMLBase.AbstractDiscreteProblem,
+            SciMLBase.AbstractODEProblem,
+            SciMLBase.AbstractDAEProblem,
+            SciMLBase.AbstractDDEProblem,
+            SciMLBase.AbstractSDEProblem,
+            SciMLBase.AbstractSDDEProblem,
+            SciMLBase.AbstractRODEProblem
+        },
+        alg, sensealg::EnzymeAdjoint,
+        u0, p, originator::SciMLBase.ADOriginator,
+        args...; kwargs...)
+    kwargs_filtered = NamedTuple(filter(x -> x[1] != :sensealg, kwargs))
+    du0 = make_zero(u0)
+    dp = make_zero(p)
+    mode = sensealg.mode
+
+    f = (u0, p) -> solve(prob, alg, args...; u0 = u0, p = p,
+            sensealg = SensitivityADPassThrough(),
+            kwargs_filtered...)
+
+    splitmode = if mode isa Forward
+        error("EnzymeAdjoint currently only allows mode=Reverse. File an issue if this is necessary.")
+    elseif mode === nothing || mode === Reverse
+        ReverseSplitWithPrimal
+    end
+
+    forward, reverse = autodiff_thunk(splitmode, Const{typeof(f)}, Duplicated, Duplicated{typeof(u0)}, Duplicated{typeof(p)})
+    tape, result, shadow_result = forward(Const(f), Duplicated(u0, du0), Duplicated(p, dp))
+
+    function enzyme_sensitivity_backpass(Δ)
+        reverse(Const(f), Duplicated(u0, du0), Duplicated(p, dp), Δ, tape)
+        if originator isa SciMLBase.TrackerOriginator ||
+           originator isa SciMLBase.ReverseDiffOriginator
+            (NoTangent(), NoTangent(), du0, dp, NoTangent(),
+                ntuple(_ -> NoTangent(), length(args))...)
+        else
+            (NoTangent(), NoTangent(), NoTangent(), du0, dp, NoTangent(),
+                ntuple(_ -> NoTangent(), length(args))...)
+        end
+    end
+    sol, enzyme_sensitivity_backpass
+end
+
+# NOTE: This is needed to prevent a method ambiguity error
+function DiffEqBase._concrete_solve_adjoint(
+        prob::AbstractNonlinearProblem, alg, sensealg::EnzymeAdjoint,
+        u0, p, originator::SciMLBase.ADOriginator,
+        args...; kwargs...)
+    kwargs_filtered = NamedTuple(filter(x -> x[1] != :sensealg, kwargs))
+
+    du0 = make_zero(u0)
+    dp = make_zero(p)
+    mode = sensealg.mode
+
+    f = (u0, p) -> solve(prob, alg, args...; u0 = u0, p = p,
+            sensealg = SensitivityADPassThrough(),
+            kwargs_filtered...)
+
+    splitmode = if mode isa Forward
+        error("EnzymeAdjoint currently only allows mode=Reverse. File an issue if this is necessary.")
+    elseif mode === nothing || mode === Reverse
+        ReverseSplitWithPrimal
+    end
+
+    forward, reverse = autodiff_thunk(splitmode, Const{typeof(f)}, Duplicated, Duplicated{typeof(u0)}, Duplicated{typeof(p)})
+    tape, result, shadow_result = forward(Const(f), Duplicated(u0, du0), Duplicated(p, dp))
+
+    function enzyme_sensitivity_backpass(Δ)
+        reverse(Const(f), Duplicated(u0, du0), Duplicated(p, dp), Δ, tape)
+        if originator isa SciMLBase.TrackerOriginator ||
+           originator isa SciMLBase.ReverseDiffOriginator
+            (NoTangent(), NoTangent(), du0, dp, NoTangent(),
+                ntuple(_ -> NoTangent(), length(args))...)
+        else
+            (NoTangent(), NoTangent(), NoTangent(), du0, dp, NoTangent(),
+                ntuple(_ -> NoTangent(), length(args))...)
+        end
+    end
+    sol, enzyme_sensitivity_backpass
+end
+
 const ENZYME_TRACKED_REAL_ERROR_MESSAGE = """
                                              `Enzyme` is not compatible with `ReverseDiffAdjoint` nor with `TrackerAdjoint`.
                                              Either choose a different adjoint method like `GaussAdjoint`,

src/sensitivity_algorithms.jl

@@ -648,6 +648,34 @@ Currently fails on almost every solver.
 """
 struct ZygoteAdjoint <: AbstractAdjointSensitivityAlgorithm{nothing, true, nothing} end
 
+"""
+EnzymeAdjoint <: AbstractAdjointSensitivityAlgorithm{nothing,true,nothing}
+
+An implementation of discrete adjoint sensitivity analysis
+using the Enzyme.jl source-to-source AD directly on the differential equation
+solver.
+
+## Constructor
+
+```julia
+EnzymeAdjoint(mode = nothing)
+```
+
+## Arugments
+
+* `mode::M` determines the autodiff mode (forward or reverse). It can be:
+  + an object subtyping `EnzymeCore.Mode` (like `EnzymeCore.Forward` or `EnzymeCore.Reverse`) if a specific mode is required
+  + `nothing` to choose the best mode automatically
+
+## SciMLProblem Support
+
+Currently fails on almost every solver.
+"""
+struct EnzymeAdjoint{M <: Union{Nothing,EnzymeCore.Mode}} <: AbstractAdjointSensitivityAlgorithm{nothing, true, nothing} 
+  mode::M
+  EnzymeAdjoint(mode = nothing) = new(mode)
+end
+
 """
 ```julia
 ForwardLSS{CS, AD, FDT, RType, gType} <: AbstractShadowingSensitivityAlgorithm{CS, AD, FDT}