fix _value_and_grad wrapper bug

Author: zuhengxu

src/optimize.jl

@@ -5,17 +5,17 @@ function pm_next!(pm, stats::NamedTuple)
     return ProgressMeter.next!(pm; showvalues=[tuple(s...) for s in pairs(stats)])
 end
 
-_wrap_in_DI_context(args...) = DifferentiationInterface.Constant.([args...]) 
+_wrap_in_DI_context(args) = DifferentiationInterface.Constant.([args...]) 
 
 function _prepare_gradient(loss, adbackend, θ, args...)
-    if isempty(args...)
+    if isempty(args)
         return DifferentiationInterface.prepare_gradient(loss, adbackend, θ)
     end
     return DifferentiationInterface.prepare_gradient(loss, adbackend, θ, _wrap_in_DI_context(args)...)
 end
 
 function _value_and_gradient(loss, prep, adbackend, θ, args...)
-    if isempty(args...)
+    if isempty(args)
         return DifferentiationInterface.value_and_gradient(loss, prep, adbackend, θ)
     end
     return DifferentiationInterface.value_and_gradient(loss, prep, adbackend, θ, _wrap_in_DI_context(args)...)

test/Project.toml

@@ -6,6 +6,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"

test/ad.jl

@@ -1,5 +1,6 @@
 @testset "DI.AD with context wrapper" begin
     f(x, y, z) = sum(abs2, x .+ y .+ z)
+    T = Float32
 
     @testset "$T" for T in [Float32, Float64]
         x = randn(T, 10)
@@ -11,9 +12,10 @@
             ADTypes.AutoZygote(),
             ADTypes.AutoForwardDiff(; chunksize=chunksize),
             ADTypes.AutoForwardDiff(),
-            ADTypes.AutoReverseDiff(false),
+            ADTypes.AutoReverseDiff(; false),
             ADTypes.AutoMooncake(; config=Mooncake.Config()),
         ]
+            at = ADTypes.AutoMooncake(; config=Mooncake.Config())
             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)
@@ -26,15 +28,18 @@ end
     @testset "$at" for at in [
         ADTypes.AutoZygote(),
         ADTypes.AutoForwardDiff(),
-        ADTypes.AutoReverseDiff(false),
+        ADTypes.AutoReverseDiff(; false),
         ADTypes.AutoMooncake(; config=Mooncake.Config()),
     ]
         @testset "$T" for T in [Float32, Float64]
             μ = 10 * ones(T, 2)
             Σ = Diagonal(4 * ones(T, 2))
             target = MvNormal(μ, Σ)
             logp(z) = logpdf(target, z)
-
+            
+            # necessary for Zygote/mooncake to differentiate through the flow
+            # prevent opt q0
+            @leaf MvNormal 
             q₀ = MvNormal(zeros(T, 2), ones(T, 2))
             flow = Bijectors.transformed(q₀, Bijectors.Shift(zero.(μ)))
 

test/runtests.jl

@@ -4,7 +4,7 @@ using Bijectors, Optimisers
 using LinearAlgebra
 using Random
 using ADTypes
-import DifferentiationInterface as DI
+# import DifferentiationInterface as DI
 using ForwardDiff, Zygote, ReverseDiff, Mooncake
 using Test