Test against Enzyme

Author: wsmoses

docs/src/api.md

@@ -93,9 +93,10 @@ See the [AD guide](https://turinglang.org/docs/tutorials/docs-10-using-turing-au
 
 | Exported symbol   | Documentation                        | Description            |
 |:----------------- |:------------------------------------ |:---------------------- |
+| `AutoEnzyme`      | [`ADTypes.AutoEnzyme`](@extref)      | Enzyme.jl backend      |
 | `AutoForwardDiff` | [`ADTypes.AutoForwardDiff`](@extref) | ForwardDiff.jl backend |
-| `AutoReverseDiff` | [`ADTypes.AutoReverseDiff`](@extref) | ReverseDiff.jl backend |
 | `AutoMooncake`    | [`ADTypes.AutoMooncake`](@extref)    | Mooncake.jl backend    |
+| `AutoReverseDiff` | [`ADTypes.AutoReverseDiff`](@extref) | ReverseDiff.jl backend |
 
 ### Debugging
 

src/Turing.jl

@@ -23,7 +23,7 @@ using Printf: Printf
 using Random: Random
 using LinearAlgebra: I
 
-using ADTypes: ADTypes, AutoForwardDiff, AutoReverseDiff, AutoMooncake
+using ADTypes: ADTypes, AutoForwardDiff, AutoReverseDiff, AutoMooncake, AutoEnzyme
 
 const DEFAULT_ADTYPE = ADTypes.AutoForwardDiff()
 
@@ -121,6 +121,7 @@ export
     AutoForwardDiff,
     AutoReverseDiff,
     AutoMooncake,
+    AutoEnzyme,
     # Debugging - Turing
     setprogress!,
     # Distributions

test/ad.jl

@@ -20,6 +20,14 @@ if INCLUDE_MOONCAKE
     using Mooncake: Mooncake
 end
 
+const INCLUDE_ENZYME = !IS_PRERELEASE
+
+if INCLUDE_ENZYME
+    import Pkg
+    Pkg.add("Enzyme")
+    using Enzyme: Enzyme
+end
+
 """Element types that are always valid for a VarInfo regardless of ADType."""
 const always_valid_eltypes = (AbstractFloat, AbstractIrrational, Integer, Rational)
 
@@ -191,6 +199,10 @@ ADTYPES = [AutoForwardDiff(), AutoReverseDiff(; compile=false)]
 if INCLUDE_MOONCAKE
     push!(ADTYPES, AutoMooncake(; config=nothing))
 end
+if INCLUDE_ENZYME
+    push!(ADTYPES, AutoEnzyme(; mode = set_runtime_activity(Forward)))
+    push!(ADTYPES, AutoEnzyme(; mode = set_runtime_activity(Reverse)))
+end
 
 # Check that ADTypeCheckContext itself works as expected.
 @testset "ADTypeCheckContext" begin