From b111c50374b6c374dcaa9528b95d06ee4ee359b1 Mon Sep 17 00:00:00 2001 From: Hossein Pourbozorg Date: Mon, 21 Jul 2025 01:48:17 +0330 Subject: [PATCH 1/2] try `.=` --- src/utils.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/utils.jl b/src/utils.jl index 0527f19d..173c982a 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -9,7 +9,7 @@ function jacobian_batched( res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2)) for i in axes(xs, 1) ChainRulesCore.@ignore_derivatives z[i, :] .= one(T) - res[i, :, :] = + res[i, :, :] .= only(DifferentiationInterface.pullback(f, icnf.compute_mode.adback, xs, (z,))) ChainRulesCore.@ignore_derivatives z[i, :] .= zero(T) end @@ -27,7 +27,7 @@ function jacobian_batched( res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2)) for i in axes(xs, 1) ChainRulesCore.@ignore_derivatives z[i, :] .= one(T) - res[:, i, :] = only( + res[:, i, :] .= only( DifferentiationInterface.pushforward(f, icnf.compute_mode.adback, xs, (z,)), ) ChainRulesCore.@ignore_derivatives z[i, :] .= zero(T) From bdadc7fb5e9709289a43435df75d16c5c067fd08 Mon Sep 17 00:00:00 2001 From: Hossein Pourbozorg Date: Tue, 22 Jul 2025 00:54:07 +0330 Subject: [PATCH 2/2] test with zygote --- benchmark/benchmarks.jl | 2 ++ examples/usage.jl | 6 +++++- test/instability_tests.jl | 1 + test/regression_tests.jl | 1 + test/smoke_tests.jl | 1 + 5 files changed, 10 insertions(+), 1 deletion(-) diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 673dcba5..a751cdfc 100644 --- a/benchmark/benchmarks.jl +++ b/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, ), ), diff --git a/examples/usage.jl b/examples/usage.jl index c6838e0f..75588bd9 100644 --- a/examples/usage.jl +++ b/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 ) diff --git a/test/instability_tests.jl b/test/instability_tests.jl index e8c52728..1b06e771 100644 --- a/test/instability_tests.jl +++ b/test/instability_tests.jl @@ -31,6 +31,7 @@ Test.@testset "Instability" begin alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(; thread = Static.True()), sensealg = SciMLSensitivity.InterpolatingAdjoint(; autodiff = true, + autojacvec = ZygoteVJP(), checkpointing = true, ), ), diff --git a/test/regression_tests.jl b/test/regression_tests.jl index f0599b7f..22cc34ae 100644 --- a/test/regression_tests.jl +++ b/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, ), ), diff --git a/test/smoke_tests.jl b/test/smoke_tests.jl index 452cdc9d..3cba3715 100644 --- a/test/smoke_tests.jl +++ b/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, ), ),