From c5b64ae9cbe498df663a9e893b828ee9b6e5a77f Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 18 May 2026 15:12:18 +0100 Subject: [PATCH] Override optimize_marginal! for LogDensityFunctionWrapper Rebuild the OptimizationProblem on each call so that updated non-marginalized parameters are passed through to Optimization.jl. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/DynamicPPLMarginalLogDensitiesExt.jl | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/ext/DynamicPPLMarginalLogDensitiesExt.jl b/ext/DynamicPPLMarginalLogDensitiesExt.jl index 905ee168b..714a3be2d 100644 --- a/ext/DynamicPPLMarginalLogDensitiesExt.jl +++ b/ext/DynamicPPLMarginalLogDensitiesExt.jl @@ -14,6 +14,20 @@ function (lw::LogDensityFunctionWrapper)(x, _) return LogDensityProblems.logdensity(lw.ldf, x) end +function MarginalLogDensities.optimize_marginal!( + mld::MarginalLogDensities.MarginalLogDensity{<:LogDensityFunctionWrapper}, p2 +) + w0 = mld.u[mld.iw] + # Rebuild the optimization problem so that the current non-marginalized + # parameters are used by Optimization.jl. + prob = MarginalLogDensities.Optimization.OptimizationProblem(mld.f_opt, w0, p2) + sol = MarginalLogDensities.Optimization.solve(prob, mld.method.solver) + wopt = sol.u::typeof(w0) + objective = sol.objective::eltype(w0) + mld.u[mld.iw] .= wopt + return wopt, objective +end + """ DynamicPPL.marginalize( model::DynamicPPL.Model,