Merge pull request #1209 from SciML/enzymeadjoint

Add EnzymeAdjoint

Author: ChrisRackauckas

docs/src/manual/differential_equation_sensitivities.md

@@ -200,6 +200,8 @@ QuadratureAdjoint
 ReverseDiffAdjoint
 TrackerAdjoint
 ZygoteAdjoint
+EnzymeAdjoint
+MooncakeAdjoint
 ForwardLSS
 AdjointLSS
 NILSS

src/SciMLSensitivity.jl

@@ -96,7 +96,7 @@ export ODEForwardSensitivityFunction, ODEForwardSensitivityProblem, SensitivityF
 
 export BacksolveAdjoint, QuadratureAdjoint, GaussAdjoint, InterpolatingAdjoint,
        TrackerAdjoint, ZygoteAdjoint, ReverseDiffAdjoint, MooncakeAdjoint,
-       ForwardSensitivity, ForwardDiffSensitivity,
+       EnzymeAdjoint, ForwardSensitivity, ForwardDiffSensitivity,
        ForwardDiffOverAdjoint,
        SteadyStateAdjoint,
        ForwardLSS, AdjointLSS, NILSS, NILSAS

src/concrete_solve.jl

@@ -1260,6 +1260,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 = Enzyme.make_zero(u0)
+    dp = Enzyme.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 Enzyme.ForwardMode
+        error("EnzymeAdjoint currently only allows mode=Reverse. File an issue if this is necessary.")
+    elseif mode === nothing || mode isa Enzyme.ReverseMode
+        Enzyme.set_runtime_activity(Enzyme.ReverseSplitWithPrimal)
+    end
+
+    forward, reverse = Enzyme.autodiff_thunk(splitmode, Enzyme.Const{typeof(f)}, Enzyme.Duplicated, Enzyme.Duplicated{typeof(u0)}, Enzyme.Duplicated{typeof(p)})
+    tape, result, shadow_result = forward(Enzyme.Const(f), Enzyme.Duplicated(copy(u0), du0), Enzyme.Duplicated(copy(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

@@ -615,6 +615,10 @@ MooncakeAdjoint <: AbstractAdjointSensitivityAlgorithm{nothing, true, nothing}
 An implementation of discrete adjoint sensitivity analysis
 using the Mooncake.jl direct differentiation. 
 
+!!! warn
+      This is currently experimental and supports only explicit solvers. It will
+      support all solvers in the future.
+
 ## Constructor
 
 ```julia
@@ -656,6 +660,9 @@ An implementation of discrete adjoint sensitivity analysis
 using the Zygote.jl source-to-source AD directly on the differential equation
 solver.
 
+!!! warn
+      This is only supports SimpleDiffEq.jl solvers due to limitations of Enzyme.
+
 ## Constructor
 
 ```julia
@@ -668,6 +675,38 @@ 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.
+
+!!! warn
+      This is currently experimental and supports only explicit solvers. It will
+      support all solvers in the future.
+
+## 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,Enzyme.EnzymeCore.Mode}} <: AbstractAdjointSensitivityAlgorithm{nothing, true, nothing} 
+  mode::M
+  EnzymeAdjoint(mode = nothing) = new{typeof(mode)}(mode)
+end
+
 """
 ```julia
 ForwardLSS{CS, AD, FDT, RType, gType} <: AbstractShadowingSensitivityAlgorithm{CS, AD, FDT}

test/concrete_solve_derivatives.jl

@@ -1,6 +1,6 @@
 using SciMLSensitivity, OrdinaryDiffEq, Zygote
 using Test, ForwardDiff
-import Tracker, ReverseDiff, ChainRulesCore, Mooncake
+import Tracker, ReverseDiff, ChainRulesCore, Mooncake, Enzyme
 
 function fiip(du, u, p, t)
     du[1] = dx = p[1] * u[1] - p[2] * u[1] * u[2]
@@ -79,6 +79,13 @@ du06, dp6 = Zygote.gradient(
         sensealg = MooncakeAdjoint())),
     u0,
     p)
+@test_broken du08, dp8 = Zygote.gradient(
+    (u0, p) -> sum(solve(prob, Tsit5(), u0 = u0, p = p,
+        abstol = 1e-14, reltol = 1e-14,
+        saveat = 0.1,
+        sensealg = EnzymeAdjoint())),
+    u0,
+    p)
 
 @test ū0≈du01 rtol=1e-12
 @test ū0 == du02
@@ -87,13 +94,15 @@ du06, dp6 = Zygote.gradient(
 #@test ū0 ≈ du05 rtol=1e-12
 @test ū0≈du06 rtol=1e-12
 @test_broken ū0≈du07 rtol=1e-12
+@test_broken ū0≈du08 rtol=1e-12
 @test adj≈dp1' rtol=1e-12
 @test adj == dp2'
 @test adj≈dp3' rtol=1e-12
 @test adj≈dp4' rtol=1e-12
 #@test adj ≈ dp5' rtol=1e-12
 @test adj≈dp6' rtol=1e-12
 @test_broken adj≈dp7' rtol=1e-12
+@test_broken adj≈dp8' rtol=1e-12
 
 ###
 ### Direct from prob
@@ -322,7 +331,14 @@ du06, dp6 = Zygote.gradient(
         sensealg = MooncakeAdjoint())),
     u0,
     p)
-du08, dp8 = Zygote.gradient(
+@test_broken du08, dp8 = Zygote.gradient(
+    (u0, p) -> sum(solve(proboop, Tsit5(), u0 = u0, p = p,
+        abstol = 1e-14, reltol = 1e-14,
+        saveat = 0.1,
+        sensealg = EnzymeAdjoint())),
+    u0,
+    p)
+du09, dp9 = Zygote.gradient(
     (u0, p) -> sum(solve(proboop, Tsit5(), u0 = u0, p = p,
         abstol = 1e-14, reltol = 1e-14,
         saveat = 0.1,
@@ -337,15 +353,17 @@ du08, dp8 = Zygote.gradient(
 #@test ū0 ≈ du05 rtol=1e-12
 @test ū0≈du06 rtol=1e-12
 @test_broken ū0≈du07 rtol=1e-12
-@test ū0≈du08 rtol=1e-12
+@test_broken ū0≈du08 rtol=1e-12
+@test ū0≈du09 rtol=1e-12
 @test adj≈dp1' rtol=1e-12
 @test adj≈dp2' rtol=1e-12
 @test adj≈dp3' rtol=1e-12
 @test adj≈dp4' rtol=1e-12
 #@test adj ≈ dp5' rtol=1e-12
 @test adj≈dp6' rtol=1e-12
 @test_broken adj≈dp7' rtol=1e-12
-@test adj≈dp8' rtol=1e-12
+@test_broken adj≈dp8' rtol=1e-12
+@test adj≈dp9' rtol=1e-12
 
 ###
 ### forward