fix bug in calculating ELBO with subsampling

Author: Red-Portal

src/algorithms/paramspacesgd/subsampledobj.jl

@@ -49,11 +49,11 @@ function estimate_objective(
 )
     (; objective, subsampling) = subobj
     sub_st = init(rng, subsampling)
-    return mean(1:length(subsampling)) do _
+    return mapreduce(+, 1:length(subsampling)) do _
         batch, sub_st, _ = step(rng, subsampling, sub_st)
         prob_sub = subsample(prob, batch)
         q_sub = subsample(q, batch)
-        estimate_objective(rng, objective, q_sub, prob_sub; kwargs...)
+        estimate_objective(rng, objective, q_sub, prob_sub; kwargs...) / length(subsampling)
     end
 end
 

test/algorithms/paramspacesgd/subsampledobj.jl

@@ -63,9 +63,8 @@ end
 
     @testset "estimate_objective batchsize=$(batchsize)" for batchsize in [1, 3, 4]
         sub_obj′ = SubsampledObjective(full_obj, batchsize, 1:n_data)
-        full_objval = estimate_objective(full_obj, q0, prob; n_samples=10^6)
-        sub_objval = estimate_objective(sub_obj′, q0, prob; n_samples=10^6)
-        @info("", full_objval, sub_objval)
+        full_objval = estimate_objective(full_obj, q0, prob; n_samples=10^8)
+        sub_objval = estimate_objective(sub_obj′, q0, prob; n_samples=10^8)
         @test full_objval ≈ sub_objval rtol=0.1
     end