test via zygote

Author: prbzrg

benchmark/Project.toml

@@ -5,6 +5,7 @@ ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
@@ -15,6 +16,7 @@ ComponentArrays = "0.15"
 DifferentiationInterface = "0.6"
 Lux = "1"
 PkgBenchmark = "0.2"
+SciMLSensitivity = "7"
 StableRNGs = "1"
 Zygote = "0.6, 0.7"
 julia = "1.10"

benchmark/benchmarks.jl

@@ -38,6 +38,12 @@ icnf = ContinuousNormalizingFlows.construct(
     steer_rate = 1.0f-1,
     λ₃ = 1.0f-2,
     rng,
+    sol_kwargs = (
+        sensealg = SciMLSensitivity.InterpolatingAdjoint(;
+            autodiff = true,
+            autojacvec = SciMLSensitivity.ZygoteVJP(),
+        ),
+    ),
 )
 ps, st = Lux.setup(icnf.rng, icnf)
 ps = ComponentArrays.ComponentArray(ps)
@@ -84,6 +90,12 @@ icnf2 = ContinuousNormalizingFlows.construct(
     steer_rate = 1.0f-1,
     λ₃ = 1.0f-2,
     rng,
+    sol_kwargs = (
+        sensealg = SciMLSensitivity.InterpolatingAdjoint(;
+            autodiff = true,
+            autojacvec = SciMLSensitivity.ZygoteVJP(),
+        ),
+    ),
 )
 
 function diff_loss_tn2(x)

test/Project.toml

@@ -12,6 +12,7 @@ Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
@@ -30,6 +31,7 @@ JET = "0.9"
 Lux = "1"
 MLJBase = "1"
 SciMLBase = "2"
+SciMLSensitivity = "7"
 StableRNGs = "1"
 TerminalLoggers = "0.1"
 Zygote = "0.6, 0.7"

test/call_tests.jl

@@ -150,6 +150,12 @@ Test.@testset "Call Tests" begin
                 resource,
                 steer_rate = convert(data_type, 1.0e-1),
                 λ₃ = convert(data_type, 1.0e-2),
+                sol_kwargs = (
+                    sensealg = SciMLSensitivity.InterpolatingAdjoint(;
+                        autodiff = true,
+                        autojacvec = SciMLSensitivity.ZygoteVJP(),
+                    ),
+                ),
             ),
             ContinuousNormalizingFlows.construct(
                 mt,
@@ -159,6 +165,12 @@ Test.@testset "Call Tests" begin
                 compute_mode,
                 inplace,
                 resource,
+                sol_kwargs = (
+                    sensealg = SciMLSensitivity.InterpolatingAdjoint(;
+                        autodiff = true,
+                        autojacvec = SciMLSensitivity.ZygoteVJP(),
+                    ),
+                ),
             ),
         )
         ps, st = Lux.setup(icnf.rng, icnf)

test/fit_tests.jl

@@ -145,6 +145,12 @@ Test.@testset "Fit Tests" begin
                 resource,
                 steer_rate = convert(data_type, 1.0e-1),
                 λ₃ = convert(data_type, 1.0e-2),
+                sol_kwargs = (
+                    sensealg = SciMLSensitivity.InterpolatingAdjoint(;
+                        autodiff = true,
+                        autojacvec = SciMLSensitivity.ZygoteVJP(),
+                    ),
+                ),
             ),
             ContinuousNormalizingFlows.construct(
                 mt,
@@ -154,6 +160,12 @@ Test.@testset "Fit Tests" begin
                 compute_mode,
                 inplace,
                 resource,
+                sol_kwargs = (
+                    sensealg = SciMLSensitivity.InterpolatingAdjoint(;
+                        autodiff = true,
+                        autojacvec = SciMLSensitivity.ZygoteVJP(),
+                    ),
+                ),
             ),
         )
         if mt <: Union{

test/instability_tests.jl

@@ -20,6 +20,12 @@ Test.@testset "Instability" begin
         tspan = (0.0f0, 13.0f0),
         steer_rate = 1.0f-1,
         λ₃ = 1.0f-2,
+        sol_kwargs = (
+            sensealg = SciMLSensitivity.InterpolatingAdjoint(;
+                autodiff = true,
+                autojacvec = SciMLSensitivity.ZygoteVJP(),
+            ),
+        ),
     )
     ps, st = Lux.setup(icnf.rng, icnf)
     ps = ComponentArrays.ComponentArray(ps)

test/regression_tests.jl

@@ -16,6 +16,12 @@ Test.@testset "Regression Tests" begin
         steer_rate = 1.0f-1,
         λ₃ = 1.0f-2,
         rng,
+        sol_kwargs = (
+            sensealg = SciMLSensitivity.InterpolatingAdjoint(;
+                autodiff = true,
+                autojacvec = SciMLSensitivity.ZygoteVJP(),
+            ),
+        ),
     )
     ps, st = Lux.setup(icnf.rng, icnf)
     ps = ComponentArrays.ComponentArray(ps)

test/runtests.jl

@@ -11,6 +11,7 @@ import ADTypes,
     Lux,
     MLJBase,
     SciMLBase,
+    SciMLSensitivity,
     StableRNGs,
     TerminalLoggers,
     Test,