switch to GaussAdjoint

Author: prbzrg

benchmark/benchmarks.jl

@@ -41,7 +41,7 @@ icnf = ContinuousNormalizingFlows.construct(
     sol_kwargs = (;
         save_everystep = false,
         alg = OrdinaryDiffEqDefault.DefaultODEAlgorithm(),
-        sensealg = SciMLSensitivity.InterpolatingAdjoint(),
+        sensealg = SciMLSensitivity.GaussAdjoint(),
     ),
 )
 
@@ -61,7 +61,7 @@ icnf2 = ContinuousNormalizingFlows.construct(
     sol_kwargs = (;
         save_everystep = false,
         alg = OrdinaryDiffEqDefault.DefaultODEAlgorithm(),
-        sensealg = SciMLSensitivity.InterpolatingAdjoint(),
+        sensealg = SciMLSensitivity.GaussAdjoint(),
     ),
 )
 

examples/usage.jl

@@ -43,7 +43,7 @@ icnf = construct(
     sol_kwargs = (;
         save_everystep = false,
         alg = DefaultODEAlgorithm(),
-        sensealg = InterpolatingAdjoint(),
+        sensealg = GaussAdjoint(),
     ), # pass to the solver
 )
 

test/ci_tests/regression_tests.jl

@@ -27,7 +27,7 @@ Test.@testset verbose = true showtiming = true failfast = false "Regression Test
         sol_kwargs = (;
             save_everystep = false,
             alg = OrdinaryDiffEqDefault.DefaultODEAlgorithm(),
-            sensealg = SciMLSensitivity.InterpolatingAdjoint(),
+            sensealg = SciMLSensitivity.GaussAdjoint(),
         ),
     )
 

test/ci_tests/smoke_tests.jl

@@ -134,7 +134,7 @@ Test.@testset verbose = true showtiming = true failfast = false "Smoke Tests" be
             sol_kwargs = (;
                 save_everystep = false,
                 alg = OrdinaryDiffEqDefault.DefaultODEAlgorithm(),
-                sensealg = SciMLSensitivity.InterpolatingAdjoint(),
+                sensealg = SciMLSensitivity.GaussAdjoint(),
             ),
         )
         ps, st = LuxCore.setup(icnf.rng, icnf)
@@ -207,7 +207,6 @@ Test.@testset verbose = true showtiming = true failfast = false "Smoke Tests" be
 
         Test.@testset verbose = true showtiming = true failfast = false "$adtype on loss" for adtype in
                                                                                               adtypes
-
             Test.@test !isnothing(DifferentiationInterface.gradient(diff_loss, adtype, ps)) broken =
                 compute_mode.adback isa ADTypes.AutoEnzyme{<:Enzyme.ForwardMode} && (
                     omode isa ContinuousNormalizingFlows.TrainMode || (

test/ci_tests/speed_tests.jl

@@ -32,7 +32,6 @@ Test.@testset verbose = true showtiming = true failfast = false "Speed Tests" be
 
     Test.@testset verbose = true showtiming = true failfast = false "$compute_mode" for compute_mode in
                                                                                         compute_modes
-
         @show compute_mode
 
         rng = StableRNGs.StableRNG(1)
@@ -63,7 +62,7 @@ Test.@testset verbose = true showtiming = true failfast = false "Speed Tests" be
             sol_kwargs = (;
                 save_everystep = false,
                 alg = OrdinaryDiffEqDefault.DefaultODEAlgorithm(),
-                sensealg = SciMLSensitivity.InterpolatingAdjoint(),
+                sensealg = SciMLSensitivity.GaussAdjoint(),
             ),
         )
 

test/quality_tests/checkby_JET_tests.jl

@@ -117,7 +117,7 @@ Test.@testset verbose = true showtiming = true failfast = false "CheckByJET" beg
             sol_kwargs = (;
                 save_everystep = false,
                 alg = OrdinaryDiffEqDefault.DefaultODEAlgorithm(),
-                sensealg = SciMLSensitivity.InterpolatingAdjoint(),
+                sensealg = SciMLSensitivity.GaussAdjoint(),
             ),
         )
         ps, st = LuxCore.setup(icnf.rng, icnf)