fixing test bug

Author: zuhengxu

Project.toml

@@ -37,7 +37,7 @@ Requires = "1"
 ReverseDiff = "1.14"
 StatsBase = "0.33, 0.34"
 Zygote = "0.6, 0.7"
-julia = "1.6"
+julia = "1.10"
 
 [extras]
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"

test/ad.jl

@@ -14,7 +14,8 @@
             ADTypes.AutoReverseDiff(false),
             ADTypes.AutoMooncake(; config=Mooncake.Config()),
         ]
-            value, grad = NormalizingFlows._value_and_gradient(f, at, x, y, z)
+            prep = NormalizingFlows._prepare_gradient(f, at, x, y, z)
+            value, grad = NormalizingFlows._value_and_gradient(f, prep, at, x, y, z)
             @test DiffResults.value(out) ≈ f(x, y, z)
             @test DiffResults.gradient(out) ≈ 2 * (x .+ y .+ z)
         end
@@ -42,8 +43,9 @@ end
 
             # check grad computation for elbo
             loss(θ, args...) = -NormalizingFlows.elbo(re(θ), args...)
+            prep = NormalizingFlows._prepare_gradient(loss, at, θ, logp, randn(T, 2, sample_per_iter))
             value, grad = NormalizingFlows._value_and_gradient(
-                loss, at, θ, logp, randn(T, 2, sample_per_iter)
+                loss, prep, at, θ, logp, randn(T, 2, sample_per_iter)
             )
 
             @test !isnothing(value)