add projection for Bijectors with MvLocationScale

Author: Red-Portal

ext/AdvancedVIBijectorsExt.jl

@@ -4,13 +4,37 @@ module AdvancedVIBijectorsExt
 if isdefined(Base, :get_extension)
     using AdvancedVI
     using Bijectors
+    using LinearAlgebra
+    using Optimisers
     using Random
 else
     using ..AdvancedVI
     using ..Bijectors
+    using ..LinearAlgebra
+    using ..Optimisers
     using ..Random
 end
 
+function AdvancedVI.update_variational_params!(
+    ::Type{<:Bijectors.TransformedDistribution{<:AdvancedVI.MvLocationScale}},
+    opt_st,
+    params,
+    restructure,
+    grad
+)
+    opt_st, params = Optimisers.update!(opt_st, params, grad)
+    q = restructure(params)
+    ϵ = q.dist.scale_eps
+
+    # Project the scale matrix to the set of positive definite triangular matrices
+    diag_idx = diagind(q.dist.scale)
+    @. q.dist.scale[diag_idx] = max(q.dist.scale[diag_idx], ϵ)
+
+    params, _ = Optimisers.destructure(q)
+
+    opt_st, params
+end
+
 function AdvancedVI.reparam_with_entropy(
     rng      ::Random.AbstractRNG,
     q        ::Bijectors.TransformedDistribution,