Test Enzyme and reexport ADTypes.AutoEnzyme

src/essential/Essential.jl

@@ -43,6 +43,7 @@ export  @model,
         setadbackend,
         setadsafe,
         ForwardDiffAD,
+        EnzymeAD,
         TrackerAD,
         ZygoteAD,
         ReverseDiffAD,

src/essential/ad.jl

@@ -12,6 +12,9 @@ end
 function _setadbackend(::Val{:forwarddiff})
     ADBACKEND[] = :forwarddiff
 end
+function _setadbackend(::Val{:enzyme})
+    ADBACKEND[] = :enzyme
+end
 function _setadbackend(::Val{:tracker})
     ADBACKEND[] = :tracker
 end
@@ -47,6 +50,7 @@ getchunksize(::ForwardDiffAD{chunk}) where chunk = chunk
 standardtag(::ForwardDiffAD{<:Any,true}) = true
 standardtag(::ForwardDiffAD) = false
 
+struct EnzymeAD <: ADBackend end
 struct TrackerAD <: ADBackend end
 struct ZygoteAD <: ADBackend end
 
@@ -64,6 +68,7 @@ ADBackend() = ADBackend(ADBACKEND[])
 ADBackend(T::Symbol) = ADBackend(Val(T))
 
 ADBackend(::Val{:forwarddiff}) = ForwardDiffAD{CHUNKSIZE[]}
+ADBackend(::Val{:enzyme}) = EnzymeAD
 ADBackend(::Val{:tracker}) = TrackerAD
 ADBackend(::Val{:zygote}) = ZygoteAD
 ADBackend(::Val{:reversediff}) = ReverseDiffAD{getrdcache()}
@@ -102,8 +107,12 @@ function LogDensityProblemsAD.ADgradient(ad::ForwardDiffAD, ℓ::Turing.LogDensi
     return LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓ; gradientconfig=config)
 end
 
+function LogDensityProblemsAD.ADgradient(::EnzymeAD, ℓ::Turing.LogDensityFunction)
+    return LogDensityProblems.ADgradient(Val(:Enzyme), ℓ)
+end
+
 function LogDensityProblemsAD.ADgradient(::TrackerAD, ℓ::Turing.LogDensityFunction)
-    return LogDensityProblemsAD.ADgradient(Val(:Tracker), ℓ)
+    return LogDensityProblems.ADgradient(Val(:Tracker), ℓ)
 end
 
 function LogDensityProblemsAD.ADgradient(::ZygoteAD, ℓ::Turing.LogDensityFunction)