Optimize code for better performance and maintainability

bench/benchmarks.jl

@@ -1,4 +1,3 @@
-
 using ADTypes
 using AdvancedVI
 using BenchmarkTools
@@ -51,8 +50,8 @@ begin
                 ("Zygote", AutoZygote()),
                 ("ForwardDiff", AutoForwardDiff()),
                 ("ReverseDiff", AutoReverseDiff()),
-                #("Mooncake", AutoMooncake(; config=Mooncake.Config())),
-                #("Enzyme", AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)),
+                ("Mooncake", AutoMooncake(; config=Mooncake.Config())),
+                # ("Enzyme", AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)),
             ],
             (familyname, family) in [
                 ("meanfield", MeanFieldGaussian(zeros(T, d), Diagonal(ones(T, d)))),

bench/normallognormal.jl

@@ -1,4 +1,3 @@
-
 struct NormalLogNormal{MX,SX,MY,SY}
     μ_x::MX
     σ_x::SX
@@ -8,7 +7,9 @@ end
 
 function LogDensityProblems.logdensity(model::NormalLogNormal, θ)
     (; μ_x, σ_x, μ_y, Σ_y) = model
-    return logpdf(LogNormal(μ_x, σ_x), θ[1]) + logpdf(MvNormal(μ_y, Σ_y), θ[2:end])
+    log_density_x = logpdf(LogNormal(μ_x, σ_x), θ[1])
+    log_density_y = logpdf(MvNormal(μ_y, Σ_y), θ[2:end])
+    return log_density_x + log_density_y
 end
 
 function LogDensityProblems.dimension(model::NormalLogNormal)

ext/AdvancedVIBijectorsExt.jl

@@ -1,4 +1,3 @@
-
 module AdvancedVIBijectorsExt
 
 if isdefined(Base, :get_extension)
@@ -60,7 +59,7 @@ function AdvancedVI.reparam_with_entropy(
     q_unconst = q.dist
     q_unconst_stop = q_stop.dist
 
-    # Draw samples and compute entropy of the uncontrained distribution
+    # Draw samples and compute entropy of the unconstrained distribution
     unconstr_samples, unconst_entropy = AdvancedVI.reparam_with_entropy(
         rng, q_unconst, q_unconst_stop, n_samples, ent_est
     )

src/AdvancedVI.jl

@@ -1,4 +1,3 @@
-
 module AdvancedVI
 
 using Accessors

src/families/location_scale.jl

@@ -1,4 +1,3 @@
-
 """
     MvLocationScale(location, scale, dist)
 
@@ -59,7 +58,8 @@ end
 
 function Distributions.logpdf(q::MvLocationScale, z::AbstractVector{<:Real})
     (; location, scale, dist) = q
-    return sum(Base.Fix1(logpdf, dist), scale \ (z - location)) - logdet(scale)
+    z_std = scale \ (z - location)
+    return sum(Base.Fix1(logpdf, dist), z_std) - logdet(scale)
 end
 
 function Distributions.rand(q::MvLocationScale)