test enzyme with reverse mode

impICNF · Jan 12, 2025 · b9bb0ef · b9bb0ef
1 parent 5b1c25e
commit b9bb0ef
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 9 deletions.
diff --git a/README.md b/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 = DIVecJacMatrixMode(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

diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl
@@ -33,9 +33,9 @@ icnf = ContinuousNormalizingFlows.construct(
     nn,
     nvars,
     naugs;
-    compute_mode = ContinuousNormalizingFlows.DIJacVecMatrixMode(
+    compute_mode = ContinuousNormalizingFlows.DIVecJacMatrixMode(
         ADTypes.AutoEnzyme(;
-            mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+            mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
             function_annotation = Enzyme.Const,
         ),
     ),
@@ -92,9 +92,9 @@ icnf2 = ContinuousNormalizingFlows.construct(
     nvars,
     naugs;
     inplace = true,
-    compute_mode = ContinuousNormalizingFlows.DIJacVecMatrixMode(
+    compute_mode = ContinuousNormalizingFlows.DIVecJacMatrixMode(
         ADTypes.AutoEnzyme(;
-            mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+            mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
             function_annotation = Enzyme.Const,
         ),
     ),

diff --git a/src/base_icnf.jl b/src/base_icnf.jl
@@ -4,9 +4,9 @@ function construct(
     nvars::Int,
     naugmented::Int = 0;
     data_type::Type{<:AbstractFloat} = Float32,
-    compute_mode::ComputeMode = DIJacVecMatrixMode(
+    compute_mode::ComputeMode = DIVecJacMatrixMode(
         ADTypes.AutoEnzyme(;
-            mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+            mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
             function_annotation = Enzyme.Const,
         ),
     ),

diff --git a/test/regression_tests.jl b/test/regression_tests.jl
@@ -11,9 +11,9 @@ Test.@testset "Regression Tests" begin
         nn,
         nvars,
         naugs;
-        compute_mode = ContinuousNormalizingFlows.DIJacVecMatrixMode(
+        compute_mode = ContinuousNormalizingFlows.DIVecJacMatrixMode(
             ADTypes.AutoEnzyme(;
-                mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+                mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
                 function_annotation = Enzyme.Const,
             ),
         ),