add ad test for nsf

Author: zuhengxu

test/ad.jl

@@ -120,3 +120,50 @@ end
         end
     end
 end
+
+@testset "AD for ELBO on NSF" begin
+    @testset "$at" for at in [
+        ADTypes.AutoZygote(),
+        ADTypes.AutoForwardDiff(),
+        ADTypes.AutoReverseDiff(; compile=false),
+        ADTypes.AutoEnzyme(;
+            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
+            function_annotation=Enzyme.Const,
+        ),
+        # it doesn't work with mooncake yet
+        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 updating params of q0
+            @leaf MvNormal
+            q₀ = MvNormal(zeros(T, 2), ones(T, 2))
+            flow = realnvp(q₀, [8, 8], 3; paramtype=T)
+
+            θ, re = Optimisers.destructure(flow)
+
+            # check grad computation for elbo
+            function loss(θ, rng, logp, sample_per_iter)
+                return -NormalizingFlows.elbo_batch(rng, re(θ), logp, sample_per_iter)
+            end
+
+            rng = Random.default_rng()
+            sample_per_iter = 10
+
+            prep = NormalizingFlows._prepare_gradient(
+                loss, at, θ, rng, logp, sample_per_iter
+            )
+            value, grad = NormalizingFlows._value_and_gradient(
+                loss, prep, at, θ, rng, logp, sample_per_iter
+            )
+
+            @test value !== nothing
+            @test all(grad .!= nothing)
+        end
+    end
+end