From bfa9de02b84c15a8eaac00d1b4844768b4275e35 Mon Sep 17 00:00:00 2001
From: Kyurae Kim <kyrkim@seas.upenn.edu>
Date: Tue, 3 Dec 2024 14:07:53 -0500
Subject: [PATCH] add basic tests for interface tests of variational objectives

---
 test/inference/scoregradelbo_locationscale.jl |  2 +-
 test/interface/repgradelbo.jl                 | 45 ++++++++++++-------
 test/interface/scoregradelbo.jl               | 30 ++++++++++++-
 3 files changed, 59 insertions(+), 18 deletions(-)

diff --git a/test/inference/scoregradelbo_locationscale.jl b/test/inference/scoregradelbo_locationscale.jl
index b19ea85e..f6c25f6f 100644
--- a/test/inference/scoregradelbo_locationscale.jl
+++ b/test/inference/scoregradelbo_locationscale.jl
@@ -17,7 +17,7 @@ end
         Dict(:Normal => normal_meanfield, :Normal => normal_fullrank),
         n_montecarlo in [1, 10],
         (objname, objective) in Dict(:ScoreGradELBO => ScoreGradELBO(n_montecarlo)),
-        (adbackname, adtype) in AD_locationscale
+        (adbackname, adtype) in AD_scoregradelbo_locationscale
 
         seed = (0x38bef07cf9cc549d)
         rng = StableRNG(seed)
diff --git a/test/interface/repgradelbo.jl b/test/interface/repgradelbo.jl
index b34dc659..da3a59ac 100644
--- a/test/interface/repgradelbo.jl
+++ b/test/interface/repgradelbo.jl
@@ -1,5 +1,14 @@
 
-using Test
+AD_repgradelbo_interface = if TEST_GROUP == "Enzyme"
+    [AutoEnzyme()]
+else
+    [
+        AutoForwardDiff(),
+        AutoReverseDiff(),
+        AutoZygote(),
+        AutoMooncake(; config=Mooncake.Config()),
+    ]
+end
 
 @testset "interface RepGradELBO" begin
     seed = (0x38bef07cf9cc549d)
@@ -9,7 +18,24 @@ using Test
 
     (; model, μ_true, L_true, n_dims, is_meanfield) = modelstats
 
-    q0 = TuringDiagMvNormal(zeros(Float64, n_dims), ones(Float64, n_dims))
+    q0 = MeanFieldGaussian(zeros(n_dims), Diagonal(ones(n_dims)))
+
+    @testset "basic" begin
+        @testset for adtype in AD_repgradelbo_interface, n_montecarlo in [1, 10]
+            obj = RepGradELBO(n_montecarlo)
+            _, _, stats, _ = optimize(
+                rng,
+                model,
+                obj,
+                q0,
+                10;
+                optimizer=Descent(1e-5),
+                show_progress=false,
+                adtype=adtype,
+            )
+            @assert isfinite(last(stats).elbo)
+        end
+    end
 
     obj = RepGradELBO(10)
     rng = StableRNG(seed)
@@ -27,17 +53,6 @@ using Test
     end
 end
 
-AD_repgradelbo_stl = if TEST_GROUP == "Enzyme"
-    [AutoEnzyme()]
-else
-    [
-        AutoForwardDiff(),
-        AutoReverseDiff(),
-        AutoZygote(),
-        AutoMooncake(; config=Mooncake.Config()),
-    ]
-end
-
 @testset "interface RepGradELBO STL variance reduction" begin
     seed = (0x38bef07cf9cc549d)
     rng = StableRNG(seed)
@@ -45,12 +60,12 @@ end
     modelstats = normal_meanfield(rng, Float64)
     (; model, μ_true, L_true, n_dims, is_meanfield) = modelstats
 
-    @testset for adtype in AD_repgradelbo_stl
+    @testset for adtype in AD_repgradelbo_interface, n_montecarlo in [1, 10]
         q_true = MeanFieldGaussian(
             Vector{eltype(μ_true)}(μ_true), Diagonal(Vector{eltype(L_true)}(diag(L_true)))
         )
         params, re = Optimisers.destructure(q_true)
-        obj = RepGradELBO(10; entropy=StickingTheLandingEntropy())
+        obj = RepGradELBO(n_montecarlo; entropy=StickingTheLandingEntropy())
         out = DiffResults.DiffResult(zero(eltype(params)), similar(params))
 
         aux = (
diff --git a/test/interface/scoregradelbo.jl b/test/interface/scoregradelbo.jl
index 8b0a3428..f368626e 100644
--- a/test/interface/scoregradelbo.jl
+++ b/test/interface/scoregradelbo.jl
@@ -1,5 +1,14 @@
 
-using Test
+AD_scoregradelbo_interface = if TEST_GROUP == "Enzyme"
+    [AutoEnzyme()]
+else
+    [
+        AutoForwardDiff(),
+        AutoReverseDiff(),
+        AutoZygote(),
+        AutoMooncake(; config=Mooncake.Config()),
+    ]
+end
 
 @testset "interface ScoreGradELBO" begin
     seed = (0x38bef07cf9cc549d)
@@ -9,7 +18,24 @@ using Test
 
     (; model, μ_true, L_true, n_dims, is_meanfield) = modelstats
 
-    q0 = TuringDiagMvNormal(zeros(Float64, n_dims), ones(Float64, n_dims))
+    q0 = MeanFieldGaussian(zeros(n_dims), Diagonal(ones(n_dims)))
+
+    @testset "basic" begin
+        @testset for adtype in AD_scoregradelbo_interface, n_montecarlo in [1, 10]
+            obj = ScoreGradELBO(n_montecarlo)
+            _, _, stats, _ = optimize(
+                rng,
+                model,
+                obj,
+                q0,
+                10;
+                optimizer=Descent(1e-5),
+                show_progress=false,
+                adtype=adtype,
+            )
+            @assert isfinite(last(stats).elbo)
+        end
+    end
 
     obj = ScoreGradELBO(10)
     rng = StableRNG(seed)