go by zygote

Author: prbzrg

benchmark/benchmarks.jl

@@ -52,6 +52,7 @@ icnf = ContinuousNormalizingFlows.construct(
         alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(; thread = Static.True()),
         sensealg = SciMLSensitivity.InterpolatingAdjoint(;
             autodiff = true,
+            autojacvec = ZygoteVJP(),
             checkpointing = true,
         ),
     ),
@@ -124,6 +125,7 @@ icnf2 = ContinuousNormalizingFlows.construct(
         alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(; thread = Static.True()),
         sensealg = SciMLSensitivity.InterpolatingAdjoint(;
             autodiff = true,
+            autojacvec = ZygoteVJP(),
             checkpointing = true,
         ),
     ),

examples/usage.jl

@@ -42,7 +42,11 @@ icnf = construct(
         abstol = eps(one(Float32)),
         maxiters = typemax(Int),
         alg = VCABM(; thread = True()),
-        sensealg = InterpolatingAdjoint(; autodiff = true, checkpointing = true),
+        sensealg = InterpolatingAdjoint(;
+            autodiff = true,
+            autojacvec = ZygoteVJP(),
+            checkpointing = true,
+        ),
     ), # pass to the solver
 )
 

test/checkby_JET_tests.jl

@@ -31,6 +31,7 @@ Test.@testset "CheckByJET" begin
             alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(; thread = Static.True()),
             sensealg = SciMLSensitivity.InterpolatingAdjoint(;
                 autodiff = true,
+                autojacvec = ZygoteVJP(),
                 checkpointing = true,
             ),
         ),

test/regression_tests.jl

@@ -27,6 +27,7 @@ Test.@testset "Regression Tests" begin
             alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(; thread = Static.True()),
             sensealg = SciMLSensitivity.InterpolatingAdjoint(;
                 autodiff = true,
+                autojacvec = ZygoteVJP(),
                 checkpointing = true,
             ),
         ),

test/smoke_tests.jl

@@ -130,6 +130,7 @@ Test.@testset "Smoke Tests" begin
                 alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(; thread = Static.True()),
                 sensealg = SciMLSensitivity.InterpolatingAdjoint(;
                     autodiff = true,
+                    autojacvec = ZygoteVJP(),
                     checkpointing = true,
                 ),
             ),