Test Enzyme on demo models

TuringLang · Feb 16, 2025 · 25b6c7e · 25b6c7e
1 parent e9cacd7
commit 25b6c7e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/test/ad.jl b/test/ad.jl
@@ -1,5 +1,6 @@
 using DynamicPPL: LogDensityFunction
 import Enzyme
+using EnzymeCore: set_runtime_activity, Forward
 
 @testset verbose = true "Automatic differentiation" begin
     @testset "Unsupported backends" begin
@@ -24,7 +25,8 @@ import Enzyme
                 ref_logp, ref_grad = LogDensityProblems.logdensity_and_gradient(ref_ldf, x)
 
                 @testset "$adtype" for adtype in [
-                    AutoEnzyme(),
+                    AutoEnzyme(; mode=set_runtime_activity(Forward, true)),
+                    AutoEnzyme(; mode=set_runtime_activity(Reverse, true)),
                     # AutoReverseDiff(; compile=false),
                     # AutoReverseDiff(; compile=true),
                     # AutoMooncake(; config=nothing),