Reduce memory footprint

Project.toml

@@ -49,7 +49,7 @@ LogDensityProblems = "2.1.0"
 LogDensityProblemsAD = "1.7.0"
 MCMCChains = "6.0.2"
 Optim = "1.7.2"
-Optimization = "3.21, 4"
+Optimization = "3.24, 4"
 OptimizationNLopt = "0.2, 0.3"
 OptimizationOptimJL = "0.2.1, 0.3, 0.4"
 PDMats = "0.11.26"

src/Pathfinder.jl

@@ -41,8 +41,8 @@ default_ad() = ADTypes.AutoForwardDiff()
 
 include("transducers.jl")
 include("woodbury.jl")
+include("lbfgs.jl")
 include("optimize.jl")
-include("inverse_hessian.jl")
 include("mvnormal.jl")
 include("elbo.jl")
 include("resample.jl")

src/elbo.jl

@@ -1,27 +1,33 @@
-function maximize_elbo(rng, logp, dists, ndraws, executor)
+function maximize_elbo(rng, logp, dists, ndraws, executor; save_samples::Bool=true)
+    dim = isempty(dists) ? 0 : length(first(dists))
+    Tdraws = Matrix{eltype(eltype(dists))}
+    draws = save_samples ? nothing : Tdraws(undef, (dim, ndraws))
     EE = Core.Compiler.return_type(
-        elbo_and_samples, Tuple{typeof(rng),typeof(logp),eltype(dists),Int}
+        elbo_and_samples!, Tuple{Tdraws,typeof(rng),typeof(logp),eltype(dists)}
     )
     estimates = similar(dists, EE)
     isempty(estimates) && return 0, estimates
+
     Folds.map!(estimates, dists, executor) do dist
-        return elbo_and_samples(rng, logp, dist, ndraws)
+        _draws = draws === nothing ? Tdraws(undef, (dim, ndraws)) : draws
+        return elbo_and_samples!(_draws, rng, logp, dist; save_samples)
     end
     _, iteration_opt = _findmax(estimates |> Transducers.Map(est -> est.value))
     return iteration_opt, estimates
 end
 
-function elbo_and_samples(rng, logp, dist, ndraws)
-    ϕ, logqϕ = rand_and_logpdf(rng, dist, ndraws)
+function elbo_and_samples!(ϕ, rng, logp, dist; save_samples::Bool=true)
+    ϕ, logqϕ = rand_and_logpdf!(rng, dist, ϕ)
     logpϕ = similar(logqϕ)
     logpϕ .= logp.(eachcol(ϕ))
     logr = logpϕ - logqϕ
     elbo = Statistics.mean(logr)
     elbo_se = sqrt(Statistics.var(logr) / length(logr))
-    return ELBOEstimate(elbo, elbo_se, ϕ, logpϕ, logqϕ, logr)
+    ϕ_save = save_samples ? copyto!(similar(ϕ), ϕ) : similar(ϕ, map(zero, size(ϕ)))
+    return ELBOEstimate(elbo, elbo_se, ϕ_save, logpϕ, logqϕ, logr)
 end
 
-struct ELBOEstimate{T,P,L<:AbstractVector{T}}
+struct ELBOEstimate{T,P<:AbstractMatrix{T},L<:Vector{T}}
     value::T
     std_err::T
     draws::P