add tests for scale projection operators, fix bug for meanfield

Author: Red-Portal

src/families/location_scale.jl

@@ -46,14 +46,14 @@ function (re::RestructureMeanField)(flat::AbstractVector)
     n_dims   = div(length(flat), 2)
     location = first(flat, n_dims)
     scale    = Diagonal(last(flat, n_dims))
-    MvLocationScale(location, scale, re.q.dist)
+    MvLocationScale(location, scale, re.q.dist, re.q.scale_eps)
 end
 
 function Optimisers.destructure(
     q::MvLocationScale{<:Diagonal, D, L}
 ) where {D, L}
     @unpack location, scale, dist = q
-    flat   = vcat(location, diag(scale))
+    flat = vcat(location, diag(scale))
     flat, RestructureMeanField(q)
 end
 # end

test/interface/location_scale.jl

@@ -138,3 +138,35 @@
         @test q         == re(λ)
     end
 end
+
+@testset "scale positive definite projection" begin
+    @testset "$(string(covtype)) $(realtype) $(bijector)" for
+        covtype     = [:meanfield, :fullrank],
+        realtype    = [Float32,     Float64],
+        bijector    = [nothing,    :identity]
+
+        d = 5
+        μ = zeros(realtype, d)
+        ϵ = sqrt(realtype(0.5))
+        q = if covtype == :fullrank
+            L = LowerTriangular(Matrix{realtype}(I,d,d))
+            FullRankGaussian(μ, L; scale_eps=ϵ)
+        elseif covtype == :meanfield
+            L = Diagonal(ones(realtype, d))
+            MeanFieldGaussian(μ, L; scale_eps=ϵ)
+        end
+        q_trans = if isnothing(bijector) 
+            q
+        else
+            Bijectors.TransformedDistribution(q, identity)
+        end
+        g = deepcopy(q)
+
+        λ, re   = Optimisers.destructure(q)
+        grad, _ = Optimisers.destructure(g)
+        opt_st  = Optimisers.setup(Descent(one(realtype)), λ)
+        _, λ′    = AdvancedVI.update_variational_params!(typeof(q), opt_st, λ, re, grad)
+        q′       = re(λ′)
+        @test all(diag(var(q′)) .≥ ϵ^2)
+    end
+end