test enzyme with reverse mode

Author: prbzrg

README.md

@@ -50,7 +50,7 @@ icnf = construct(
     nn,
     nvars, # number of variables
     naugs; # number of augmented dimensions
-    # compute_mode = DIJacVecMatrixMode(AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Forward), function_annotation = Enzyme.Const)), # process data in batches and use Enzyme
+    # compute_mode = DIJacVecMatrixMode(AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation = Enzyme.Const)), # process data in batches and use Enzyme
     # inplace = true, # use the inplace version of functions
     # resource = CUDALibs(), # process data by GPU
     tspan = (0.0f0, 13.0f0), # have bigger time span

benchmark/benchmarks.jl

@@ -35,7 +35,7 @@ icnf = ContinuousNormalizingFlows.construct(
     naugs;
     compute_mode = ContinuousNormalizingFlows.DIJacVecMatrixMode(
         ADTypes.AutoEnzyme(;
-            mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+            mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
             function_annotation = Enzyme.Const,
         ),
     ),
@@ -94,7 +94,7 @@ icnf2 = ContinuousNormalizingFlows.construct(
     inplace = true,
     compute_mode = ContinuousNormalizingFlows.DIJacVecMatrixMode(
         ADTypes.AutoEnzyme(;
-            mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+            mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
             function_annotation = Enzyme.Const,
         ),
     ),

src/base_icnf.jl

@@ -6,7 +6,7 @@ function construct(
     data_type::Type{<:AbstractFloat} = Float32,
     compute_mode::ComputeMode = DIJacVecMatrixMode(
         ADTypes.AutoEnzyme(;
-            mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+            mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
             function_annotation = Enzyme.Const,
         ),
     ),

test/regression_tests.jl

@@ -13,7 +13,7 @@ Test.@testset "Regression Tests" begin
         naugs;
         compute_mode = ContinuousNormalizingFlows.DIJacVecMatrixMode(
             ADTypes.AutoEnzyme(;
-                mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+                mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
                 function_annotation = Enzyme.Const,
             ),
         ),