use QuadratureAdjoint

Author: prbzrg

examples/usage.jl

@@ -56,10 +56,11 @@ icnf = ICNF(;
         reltol = 1.0e-4,
         abstol = 1.0e-8,
         alg = VCABM(; thread = Threaded()),
-        sensealg = GaussAdjoint(;
-            checkpointing = true,
+        sensealg = QuadratureAdjoint(;
             autodiff = true,
             autojacvec = ZygoteVJP(),
+            reltol = 1.0e-4,
+            abstol = 1.0e-8,
         ),
         progress = false,
         verbose = Detailed(),
@@ -84,7 +85,7 @@ if !isfile(icnf_mach_fn)
         optimizers = (
             OptimiserChain(
                 WeightDecay(; lambda = 1.0e-4),
-                ClipNorm(1.0, 2.0; throw = true),
+                ClipNorm(10.0, 2.0; throw = true),
                 Adam(; eta = 0.001, beta = (0.9, 0.999), epsilon = 1.0e-8),
             ),
         ),

src/core/icnf.jl

@@ -89,18 +89,15 @@ function ICNF(;
         alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(;
             thread = FastBroadcast.Threaded(),
         ),
-        sensealg = ifelse(
-            (!inplace && compute_mode isa LuxMatrixMode),
-            SciMLSensitivity.GaussAdjoint(;
-                checkpointing = true,
-                autodiff = true,
-                autojacvec = SciMLSensitivity.ZygoteVJP(),
-            ),
-            SciMLSensitivity.InterpolatingAdjoint(;
-                checkpointing = true,
-                autodiff = true,
-                autojacvec = true,
+        sensealg = SciMLSensitivity.QuadratureAdjoint(;
+            autodiff = true,
+            autojacvec = ifelse(
+                (!inplace && compute_mode isa LuxMatrixMode),
+                SciMLSensitivity.ZygoteVJP(),
+                true,
             ),
+            reltol = convert(data_type, 1.0e-4),
+            abstol = convert(data_type, 1.0e-8),
         ),
         progress = false,
         verbose = SciMLLogging.Detailed(),

src/exts/mlj_ext/core_cond_icnf.jl

@@ -14,7 +14,7 @@ function CondICNFModel(;
         Optimisers.OptimiserChain(
             Optimisers.WeightDecay(; lambda = convert(eltype(icnf), 1.0e-4)),
             Optimisers.ClipNorm(
-                one(eltype(icnf)),
+                convert(eltype(icnf), 10.0),
                 convert(eltype(icnf), 2.0);
                 throw = true,
             ),

src/exts/mlj_ext/core_icnf.jl

@@ -14,7 +14,7 @@ function ICNFModel(;
         Optimisers.OptimiserChain(
             Optimisers.WeightDecay(; lambda = convert(eltype(icnf), 1.0e-4)),
             Optimisers.ClipNorm(
-                one(eltype(icnf)),
+                convert(eltype(icnf), 10.0),
                 convert(eltype(icnf), 2.0);
                 throw = true,
             ),