From f689bd8324f1aaa7c089ea8343cd20bfd8751d17 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Fri, 24 Apr 2026 16:36:28 +0100 Subject: [PATCH 01/31] Migrate from DifferentiationInterface to AbstractPPL evaluator interface Replace `DifferentiationInterface` with `AbstractPPL.prepare` / `AbstractPPL.value_and_gradient` throughout. Key changes: - `_prepare_gradient` / `_value_and_gradient!` now wrap an `AbstractPPL` prepared evaluator in a `_VIGradPrep` struct that holds an `aux_ref` so auxiliary inputs can be swapped without re-preparing. - `DynamicPPLModelLogDensityFunction` stores `model_ref` and `loglikeadj_ref` as `Ref`s; `subsample` mutates them in-place instead of creating a new struct via `@set`, keeping the prepared evaluator valid across subsampling steps. - Drop second-order (Hessian) support; `use_hessian=true` now warns and is ignored. - Pin dev branches via `[sources]`: `AbstractPPL@evaluator-interface` and `DynamicPPL@adproblems-interface`. Co-Authored-By: Claude Sonnet 4.6 --- Project.toml | 10 +- ext/AdvancedVIDynamicPPLExt.jl | 149 ++++++------------ src/AdvancedVI.jl | 50 +++--- src/algorithms/abstractobjective.jl | 2 +- src/algorithms/common.jl | 2 +- src/algorithms/fisherminbatchmatch.jl | 18 +-- src/algorithms/klminnaturalgraddescent.jl | 10 +- src/algorithms/klminsqrtnaturalgraddescent.jl | 6 +- src/algorithms/klminwassfwdbwd.jl | 6 +- src/algorithms/subsampledobjective.jl | 2 +- src/optimization/rules.jl | 8 +- src/reshuffling.jl | 2 +- test/Project.toml | 6 + test/integration/dynamicppl.jl | 14 +- test/runtests.jl | 1 + 15 files changed, 129 insertions(+), 157 deletions(-) diff --git a/Project.toml b/Project.toml index 28d4a3035..886401ebe 100644 --- a/Project.toml +++ b/Project.toml @@ -4,10 +4,10 @@ version = "0.7.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" @@ -29,14 +29,14 @@ DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" AdvancedVIEnzymeExt = ["Enzyme", "ChainRulesCore"] AdvancedVIMooncakeExt = ["Mooncake", "ChainRulesCore"] AdvancedVIReverseDiffExt = ["ReverseDiff", "ChainRulesCore"] -AdvancedVIDynamicPPLExt = ["DynamicPPL", "Accessors", "Distributions", "DifferentiationInterface", "LogDensityProblems"] +AdvancedVIDynamicPPLExt = ["DynamicPPL", "Accessors", "Distributions", "LogDensityProblems"] [compat] ADTypes = "1" +AbstractPPL = "0.14" Accessors = "0.1" ChainRulesCore = "1" DiffResults = "1" -DifferentiationInterface = "0.6, 0.7" Distributions = "0.25.111" DocStringExtensions = "0.8, 0.9" DynamicPPL = "0.40, 0.41" @@ -63,3 +63,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] test = ["Pkg", "Test"] + +[sources] +AbstractPPL = {url = "https://github.com/TuringLang/AbstractPPL.jl", rev = "evaluator-interface"} +DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "adproblems-interface"} diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index eaff70ae7..08c4bcf4d 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -3,48 +3,20 @@ module AdvancedVIDynamicPPLExt using ADTypes: ADTypes using Accessors using AdvancedVI: AdvancedVI -using DifferentiationInterface: DifferentiationInterface +using AbstractPPL: AbstractPPL using Distributions: Distributions using DynamicPPL: DynamicPPL using LogDensityProblems: LogDensityProblems using Random -adtype_capabilities(::Type{Nothing}) = LogDensityProblems.LogDensityOrder{0}() +function adtype_capabilities(::Type{Nothing}) + return LogDensityProblems.LogDensityOrder{0}() +end function adtype_capabilities(::Type{<:ADTypes.AbstractADType}) return LogDensityProblems.LogDensityOrder{1}() end -function adtype_capabilities( - ::Type{ - <:Union{ - <:ADTypes.AutoForwardDiff, - <:ADTypes.AutoReverseDiff, - <:ADTypes.AutoMooncake, - <:ADTypes.AutoEnzyme, - <:DifferentiationInterface.SecondOrder, - }, - }, -) - return LogDensityProblems.LogDensityOrder{2}() -end - -struct DynamicPPLModelLogDensityFunction{ - Model<:DynamicPPL.Model, - LogLikeAdj<:Real, - VarInfo<:DynamicPPL.AbstractVarInfo, - ADType<:ADTypes.AbstractADType, - PrepGrad<:Union{Nothing,DifferentiationInterface.GradientPrep}, - PrepHess<:Union{Nothing,DifferentiationInterface.HessianPrep}, -} - model::Model - loglikeadj::LogLikeAdj - varinfo::VarInfo - adtype::ADType - prep_grad::PrepGrad - prep_hess::PrepHess -end - function logdensity_impl( params, model::DynamicPPL.Model, loglikeadj::Real, varinfo::DynamicPPL.AbstractVarInfo ) @@ -63,14 +35,31 @@ function subsample_dynamicpplmodel( return DynamicPPL.Model{Threaded}(model.f, model.args, new_kwargs, model.context) end +struct DynamicPPLModelLogDensityFunction{ + Model<:DynamicPPL.Model, + VarInfo<:DynamicPPL.AbstractVarInfo, + ADType<:Union{Nothing,ADTypes.AbstractADType}, + PrepGrad, +} + model::Model + varinfo::VarInfo + adtype::ADType + # Refs are updated in-place by subsample; the prepared AD evaluator reads + # through them on every call, so the prep remains valid across subsampling. + model_ref::Ref{Any} + loglikeadj_ref::Ref{Float64} + prep_grad::PrepGrad +end + function DynamicPPLModelLogDensityFunction( model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo; - use_hessian::Bool=true, + use_hessian::Bool=false, adtype::Union{Nothing,ADTypes.AbstractADType}=nothing, loglikeadj::Real=1.0, subsampling::Union{Nothing,AdvancedVI.AbstractSubsampling}=nothing, ) + use_hessian && @warn "`use_hessian` is no longer supported and will be ignored." model_sub = if isnothing(subsampling) model else @@ -82,92 +71,45 @@ function DynamicPPLModelLogDensityFunction( params = [val for val in varinfo[:]] cap = adtype_capabilities(typeof(adtype)) + + model_ref = Ref{Any}(model_sub) + loglikeadj_ref = Ref{Float64}(float(loglikeadj)) + prep_grad = if cap >= LogDensityProblems.LogDensityOrder{1}() - DifferentiationInterface.prepare_gradient( - logdensity_impl, - DifferentiationInterface.inner(adtype), + AbstractPPL.prepare( + adtype, + params -> logdensity_impl(params, model_ref[], loglikeadj_ref[], varinfo), params, - DifferentiationInterface.Constant(model_sub), - DifferentiationInterface.Constant(loglikeadj), - DifferentiationInterface.Constant(varinfo), ) else nothing end - prep_hess = if cap >= LogDensityProblems.LogDensityOrder{2}() && use_hessian - try - DifferentiationInterface.prepare_hessian( - logdensity_impl, - adtype, - params, - DifferentiationInterface.Constant(model_sub), - DifferentiationInterface.Constant(loglikeadj), - DifferentiationInterface.Constant(varinfo), - ) - catch - @warn "The selected AD backend has second-order capabilities but `DifferentiationInterface.prepare_hessian` failed. AdvancedVI will treat the model to only have first-order capability." - nothing - end - else - nothing - end - return DynamicPPLModelLogDensityFunction{ - typeof(model), - typeof(loglikeadj), - typeof(varinfo), - typeof(adtype), - typeof(prep_grad), - typeof(prep_hess), - }( - model, loglikeadj, varinfo, adtype, prep_grad, prep_hess + + return DynamicPPLModelLogDensityFunction( + model, varinfo, adtype, model_ref, loglikeadj_ref, prep_grad ) end function LogDensityProblems.logdensity(prob::DynamicPPLModelLogDensityFunction, params) - (; model, loglikeadj, varinfo) = prob - return logdensity_impl(params, model, loglikeadj, varinfo) + return logdensity_impl(params, prob.model_ref[], prob.loglikeadj_ref[], prob.varinfo) end function LogDensityProblems.logdensity_and_gradient( prob::DynamicPPLModelLogDensityFunction, params ) - (; model, adtype, loglikeadj, varinfo, prep_grad) = prob - return DifferentiationInterface.value_and_gradient( - logdensity_impl, - prep_grad, - DifferentiationInterface.inner(adtype), - params, - DifferentiationInterface.Constant(model), - DifferentiationInterface.Constant(loglikeadj), - DifferentiationInterface.Constant(varinfo), - ) + return AbstractPPL.value_and_gradient(prob.prep_grad, params) end -function LogDensityProblems.logdensity_gradient_and_hessian( - prob::DynamicPPLModelLogDensityFunction, params -) - (; model, adtype, loglikeadj, varinfo, prep_hess) = prob - return DifferentiationInterface.value_gradient_and_hessian( - logdensity_impl, - prep_hess, - adtype, - params, - DifferentiationInterface.Constant(model), - DifferentiationInterface.Constant(loglikeadj), - DifferentiationInterface.Constant(varinfo), - ) +function LogDensityProblems.capabilities( + ::Type{<:DynamicPPLModelLogDensityFunction{M,V,Nothing,G}} +) where {M,V,G} + return LogDensityProblems.LogDensityOrder{0}() end function LogDensityProblems.capabilities( - ::Type{<:DynamicPPLModelLogDensityFunction{M,L,V,ADType,PG,PH}} -) where {M,L,V,ADType<:ADTypes.AbstractADType,PG,PH} - return if PH != Nothing - LogDensityProblems.LogDensityOrder{2}() - elseif PG != Nothing - LogDensityProblems.LogDensityOrder{1}() - else - LogDensityProblems.LogDensityOrder{0}() - end + ::Type{<:DynamicPPLModelLogDensityFunction{M,V,<:ADTypes.AbstractADType,G}} +) where {M,V,G} + return LogDensityProblems.LogDensityOrder{1}() end function LogDensityProblems.dimension(prob::DynamicPPLModelLogDensityFunction) @@ -180,7 +122,7 @@ function AdvancedVI.subsample(prob::DynamicPPLModelLogDensityFunction, batch) if !haskey(model.defaults, :datapoints) throw( ArgumentError( - "Subsampling is turned on, but the model does not have have a `datapoints` keyword argument.", + "Subsampling is turned on, but the model does not have a `datapoints` keyword argument.", ), ) end @@ -190,9 +132,10 @@ function AdvancedVI.subsample(prob::DynamicPPLModelLogDensityFunction, batch) model_sub = subsample_dynamicpplmodel(model, batch) loglikeadj = n_datapoints / batchsize - prob′ = @set prob.model = model_sub - prob′′ = @set prob′.loglikeadj = loglikeadj - return prob′′ + prob.model_ref[] = model_sub + prob.loglikeadj_ref[] = loglikeadj + + return prob end end diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 1d07fd975..a0ce225de 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -17,13 +17,20 @@ using LogDensityProblems using ADTypes using DiffResults -using DifferentiationInterface +using AbstractPPL: AbstractPPL using ChainRulesCore: ChainRulesCore using FillArrays using StatsBase +# Holds the AbstractPPL prepared evaluator together with the aux Ref so that +# _value_and_gradient! can update aux before every evaluation. +struct _VIGradPrep{P,R} + prepared::P + aux_ref::R +end + # Derivatives """ _value_and_gradient!(f, out, ad, x, aux) @@ -33,9 +40,9 @@ Evaluate the value and gradient of a function `f` at `x` using the automatic dif `f` may receive auxiliary input as `f(x,aux)`. # Arguments -- `ad::ADTypes.AbstractADType`: +- `ad::ADTypes.AbstractADType`: automatic differentiation backend. Currently supports - `ADTypes.AutoZygote()`, `ADTypes.ForwardDiff()`, `ADTypes.ReverseDiff()`, + `ADTypes.AutoZygote()`, `ADTypes.ForwardDiff()`, `ADTypes.ReverseDiff()`, `ADTypes.AutoMooncake()` and `ADTypes.AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), @@ -45,31 +52,36 @@ Evaluate the value and gradient of a function `f` at `x` using the automatic dif - `f`: Function subject to differentiation. - `x`: The point to evaluate the gradient. - `aux`: Auxiliary input passed to `f`. -- `prep`: Output of `DifferentiationInterface.prepare_gradient`. +- `prep`: Output of `_prepare_gradient`. - `out::DiffResults.MutableDiffResult`: Buffer to contain the output gradient and function value. """ function _value_and_gradient!( f, out::DiffResults.MutableDiffResult, ad::ADTypes.AbstractADType, x, aux ) - grad_buf = DiffResults.gradient(out) - y, _ = DifferentiationInterface.value_and_gradient!(f, grad_buf, ad, x, Constant(aux)) - DiffResults.value!(out, y) + prepared = AbstractPPL.prepare(ad, Base.Fix2(f, aux), x) + val, grad = AbstractPPL.value_and_gradient(prepared, x) + DiffResults.value!(out, val) + copyto!(DiffResults.gradient(out), grad) return out end function _value_and_gradient!( - f, out::DiffResults.MutableDiffResult, prep, ad::ADTypes.AbstractADType, x, aux + f, + out::DiffResults.MutableDiffResult, + prep::_VIGradPrep, + ad::ADTypes.AbstractADType, + x, + aux, ) - grad_buf = DiffResults.gradient(out) - y, _ = DifferentiationInterface.value_and_gradient!( - f, grad_buf, prep, ad, x, Constant(aux) - ) - DiffResults.value!(out, y) + prep.aux_ref[] = aux + val, grad = AbstractPPL.value_and_gradient(prep.prepared, x) + DiffResults.value!(out, val) + copyto!(DiffResults.gradient(out), grad) return out end """ - _prepare_gradient!(f, ad, x, aux) + _prepare_gradient(f, ad, x, aux) Prepare AD backend for taking gradients of a function `f` at `x` using the automatic differentiation backend `ad`. @@ -80,7 +92,9 @@ Prepare AD backend for taking gradients of a function `f` at `x` using the autom - `aux`: Auxiliary input passed to `f`. """ function _prepare_gradient(f, ad::ADTypes.AbstractADType, x, aux) - return DifferentiationInterface.prepare_gradient(f, ad, x, Constant(aux)) + aux_ref = Ref(aux) + prepared = AbstractPPL.prepare(ad, x -> f(x, aux_ref[]), x) + return _VIGradPrep(prepared, aux_ref) end """ @@ -238,7 +252,7 @@ function step( objargs...; kwargs..., ) - nothing + return nothing end """ @@ -276,11 +290,11 @@ Please refer to the respective documentation of each algorithm for more info. function estimate_objective( ::Random.AbstractRNG, ::AbstractVariationalAlgorithm, q, prob; kwargs... ) - nothing + return nothing end function estimate_objective(alg::AbstractVariationalAlgorithm, q, prob; kwargs...) - estimate_objective(Random.default_rng(), alg, q, prob; kwargs...) + return estimate_objective(Random.default_rng(), alg, q, prob; kwargs...) end export estimate_objective diff --git a/src/algorithms/abstractobjective.jl b/src/algorithms/abstractobjective.jl index 65316c7e7..ff027bb54 100644 --- a/src/algorithms/abstractobjective.jl +++ b/src/algorithms/abstractobjective.jl @@ -31,7 +31,7 @@ function init( ::Any, ::Any, ) - nothing + return nothing end """ diff --git a/src/algorithms/common.jl b/src/algorithms/common.jl index 0b99ff0d0..96187fe6a 100644 --- a/src/algorithms/common.jl +++ b/src/algorithms/common.jl @@ -116,5 +116,5 @@ function step( ) info = !isnothing(info′) ? merge(info′, info) : info end - state, false, info + return state, false, info end diff --git a/src/algorithms/fisherminbatchmatch.jl b/src/algorithms/fisherminbatchmatch.jl index b794a12af..f3ea6fd8d 100644 --- a/src/algorithms/fisherminbatchmatch.jl +++ b/src/algorithms/fisherminbatchmatch.jl @@ -89,14 +89,14 @@ function rand_batch_match_samples_with_objective!( μ = q.location C = q.scale u = Random.randn!(rng, u_buf) - z = C*u .+ μ + z = C * u .+ μ logπ_sum = zero(eltype(μ)) for b in 1:n_samples logπb, gb = LogDensityProblems.logdensity_and_gradient(prob, view(z, :, b)) grad_buf[:, b] = gb logπ_sum += logπb end - logπ_avg = logπ_sum/n_samples + logπ_avg = logπ_sum / n_samples # Estimate objective values # @@ -105,7 +105,7 @@ function rand_batch_match_samples_with_objective!( # = E[| C' ( -(CC')\((Cu + μ) - μ) - ∇logπ(z)) |^2] (z = Cu + μ) # = E[| C' ( -(CC')\(Cu) - ∇logπ(z)) |^2] # = E[| -u - C'∇logπ(z)) |^2] - fisher = sum(abs2, -u_buf - (C'*grad_buf))/n_samples + fisher = sum(abs2, -u_buf - (C' * grad_buf)) / n_samples return u_buf, z, grad_buf, fisher, logπ_avg end @@ -145,13 +145,13 @@ function step( gbar, Γ = mean_and_cov(grad_buf, 2) μmz = μ - zbar - λ = convert(eltype(μ), d*n_samples / iteration) + λ = convert(eltype(μ), d * n_samples / iteration) - U = Symmetric(λ*Γ + (λ/(1 + λ)*gbar)*gbar') - V = Symmetric(Σ + λ*C + (λ/(1 + λ)*μmz)*μmz') + U = Symmetric(λ * Γ + (λ / (1 + λ) * gbar) * gbar') + V = Symmetric(Σ + λ * C + (λ / (1 + λ) * μmz) * μmz') - Σ′ = Hermitian(2*V/(I + real(sqrt(I + 4*U*V)))) - μ′ = 1/(1 + λ)*μ + λ/(1 + λ)*(Σ′*gbar + zbar) + Σ′ = Hermitian(2 * V / (I + real(sqrt(I + 4 * U * V)))) + μ′ = 1 / (1 + λ) * μ + λ / (1 + λ) * (Σ′ * gbar + zbar) q′ = MvLocationScale(μ′[:, 1], cholesky(Σ′).L, q.dist) elbo = logπ_avg + entropy(q) @@ -163,7 +163,7 @@ function step( info′ = callback(; rng, iteration, q, state) info = !isnothing(info′) ? merge(info′, info) : info end - state, false, info + return state, false, info end """ diff --git a/src/algorithms/klminnaturalgraddescent.jl b/src/algorithms/klminnaturalgraddescent.jl index c101537a4..596cd372c 100644 --- a/src/algorithms/klminnaturalgraddescent.jl +++ b/src/algorithms/klminnaturalgraddescent.jl @@ -81,10 +81,10 @@ function init( grad_buf = Vector{eltype(q_init.location)}(undef, n_dims) hess_buf = Matrix{eltype(q_init.location)}(undef, n_dims, n_dims) scale = q_init.scale - qcov = Hermitian(scale*scale') + qcov = Hermitian(scale * scale') scale_inv = inv(scale) prec_chol = scale_inv' - prec = Hermitian(prec_chol*prec_chol') + prec = Hermitian(prec_chol * prec_chol') return KLMinNaturalGradDescentState( q_init, prob, prec, qcov, 0, sub_st, grad_buf, hess_buf ) @@ -127,7 +127,7 @@ function step( # Handling the positive-definite constraint in the Bayesian learning rule. # In ICML 2020. G_hat = S - (-hess_buf) - Hermitian(S - η*G_hat + η^2/2*G_hat*qcov*G_hat) + Hermitian(S - η * G_hat + η^2 / 2 * G_hat * qcov * G_hat) else Hermitian(((1 - η) * S + η * (-hess_buf))) end @@ -136,7 +136,7 @@ function step( prec_chol = cholesky(S′).L prec_chol_inv = inv(prec_chol) scale = prec_chol_inv' - qcov = Hermitian(scale*scale') + qcov = Hermitian(scale * scale') q′ = MvLocationScale(m′, scale, q.dist) state = KLMinNaturalGradDescentState( @@ -149,7 +149,7 @@ function step( info′ = callback(; rng, iteration, q=q′, info) info = !isnothing(info′) ? merge(info′, info) : info end - state, false, info + return state, false, info end """ diff --git a/src/algorithms/klminsqrtnaturalgraddescent.jl b/src/algorithms/klminsqrtnaturalgraddescent.jl index 622c111c3..3b29c3b2b 100644 --- a/src/algorithms/klminsqrtnaturalgraddescent.jl +++ b/src/algorithms/klminsqrtnaturalgraddescent.jl @@ -105,8 +105,8 @@ function step( rng, q, n_samples, grad_buf, hess_buf, prob_sub ) - CtHCmI = C'*(-hess_buf)*C - I - CtHCmI_tril = LowerTriangular(tril(CtHCmI) - Diagonal(diag(CtHCmI))/2) + CtHCmI = C' * (-hess_buf) * C - I + CtHCmI_tril = LowerTriangular(tril(CtHCmI) - Diagonal(diag(CtHCmI)) / 2) m′ = m - η * C * (C' * -grad_buf) C′ = C - η * C * CtHCmI_tril @@ -123,7 +123,7 @@ function step( info′ = callback(; rng, iteration, q=q′, info) info = !isnothing(info′) ? merge(info′, info) : info end - state, false, info + return state, false, info end """ diff --git a/src/algorithms/klminwassfwdbwd.jl b/src/algorithms/klminwassfwdbwd.jl index 602f4be41..40577bfc4 100644 --- a/src/algorithms/klminwassfwdbwd.jl +++ b/src/algorithms/klminwassfwdbwd.jl @@ -104,10 +104,10 @@ function step( m′ = m - η * (-grad_buf) M = I - η * (-hess_buf') - Σ_half = Hermitian(M*Σ*M') + Σ_half = Hermitian(M * Σ * M') # Compute the JKO proximal operator - Σ′ = (Σ_half + 2*η*I + sqrt(Hermitian(Σ_half*(Σ_half + 4*η*I))))/2 + Σ′ = (Σ_half + 2 * η * I + sqrt(Hermitian(Σ_half * (Σ_half + 4 * η * I)))) / 2 q′ = MvLocationScale(m′, cholesky(Σ′).L, q.dist) state = KLMinWassFwdBwdState(q′, prob, Σ′, iteration, sub_st′, grad_buf, hess_buf) @@ -118,7 +118,7 @@ function step( info′ = callback(; rng, iteration, q=q′, info) info = !isnothing(info′) ? merge(info′, info) : info end - state, false, info + return state, false, info end """ diff --git a/src/algorithms/subsampledobjective.jl b/src/algorithms/subsampledobjective.jl index 21734f7ca..ad1f858ef 100644 --- a/src/algorithms/subsampledobjective.jl +++ b/src/algorithms/subsampledobjective.jl @@ -32,7 +32,7 @@ function init( sub_st = init(rng, subsampling) # This is necessary to ensure that `init` sees the type "conditioned" on a minibatch - # when calling `DifferentiationInterface.prepare_*` inside it. + # so that any prepared AD evaluator inside it sees the correct batch-subsampled type. batch, _, _ = step(rng, subsampling, sub_st, true) prob_sub = subsample(prob, batch) q_init_sub = subsample(q_init, batch) diff --git a/src/optimization/rules.jl b/src/optimization/rules.jl index 2025632e3..5015f4ebe 100644 --- a/src/optimization/rules.jl +++ b/src/optimization/rules.jl @@ -18,7 +18,9 @@ Optimisers.@def struct DoWG <: Optimisers.AbstractRule alpha = 1e-6 end -Optimisers.init(o::DoWG, x::AbstractArray{T}) where {T} = (copy(x), zero(T), T(o.alpha)*(1 + norm(x))) +function Optimisers.init(o::DoWG, x::AbstractArray{T}) where {T} + return (copy(x), zero(T), T(o.alpha) * (1 + norm(x))) +end function Optimisers.apply!(::DoWG, state, x::AbstractArray{T}, dx) where {T} x0, v, r = state @@ -47,7 +49,9 @@ Optimisers.@def struct DoG <: Optimisers.AbstractRule alpha = 1e-6 end -Optimisers.init(o::DoG, x::AbstractArray{T}) where {T} = (copy(x), zero(T), T(o.alpha)*(1 + norm(x))) +function Optimisers.init(o::DoG, x::AbstractArray{T}) where {T} + return (copy(x), zero(T), T(o.alpha) * (1 + norm(x))) +end function Optimisers.apply!(::DoG, state, x::AbstractArray{T}, dx) where {T} x0, v, r = state diff --git a/src/reshuffling.jl b/src/reshuffling.jl index e0e50cfbd..dd7f95aa8 100644 --- a/src/reshuffling.jl +++ b/src/reshuffling.jl @@ -49,7 +49,7 @@ function step( # Ignore the trailing batch if its size is smaller than `batchsize`. # This should only be used when estimating gradients during optimization. # This is necessary to ensure that all batches have the same size. - # Otherwise, `DifferentiationInterface.prepare_*` behaves incorrectly. + # Otherwise, prepared AD evaluators may see inconsistent batch sizes. (sub_step, batch), iterator = Iterators.peel(iterator) end epoch = epoch + 1 diff --git a/test/Project.toml b/test/Project.toml index ce8ea455b..6aad881ab 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,5 +1,7 @@ [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" +AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" @@ -23,6 +25,10 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +[sources] +AbstractPPL = {url = "https://github.com/TuringLang/AbstractPPL.jl", rev = "evaluator-interface"} +DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "adproblems-interface"} + [compat] ADTypes = "0.2.1, 1" DiffResults = "1" diff --git a/test/integration/dynamicppl.jl b/test/integration/dynamicppl.jl index 43a38d3b8..a3ed01d20 100644 --- a/test/integration/dynamicppl.jl +++ b/test/integration/dynamicppl.jl @@ -1,7 +1,7 @@ @testset "DynamicPPL" begin DynamicPPL.@model function normal(μ) - x ~ MvNormal(μ, I) + return x ~ MvNormal(μ, I) end DynamicPPL.@model function normal_subsampled(μs; datapoints=1:size(μs, 2)) @@ -22,18 +22,18 @@ alg = KLMinRepGradProxDescent(AD) d = LogDensityProblems.dimension(prob) - q0 = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6*I, d, d))) + q0 = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6 * I, d, d))) q, _, _ = AdvancedVI.optimize(alg, 1000, prob, q0; show_progress=false) Δλ0 = sum(abs2, q0.location - μ_true) Δλ = sum(abs2, q.location - μ_true) - @test Δλ ≤ Δλ0/2 + @test Δλ ≤ Δλ0 / 2 end @testset "subsampling" begin n_data = 32 - μs = 3*randn(2, n_data) - μ_true = mean(μs, dims=2)[:, 1] + μs = 3 * randn(2, n_data) + μ_true = mean(μs; dims=2)[:, 1] model = normal_subsampled(μs) vi = DynamicPPL.VarInfo(model) @@ -48,11 +48,11 @@ alg = KLMinRepGradProxDescent(AD; subsampling) d = LogDensityProblems.dimension(prob) - q0 = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6*I, d, d))) + q0 = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6 * I, d, d))) q, _, _ = AdvancedVI.optimize(alg, 1000, prob, q0; show_progress=false) Δλ0 = sum(abs2, q0.location - μ_true) Δλ = sum(abs2, q.location - μ_true) - @test Δλ ≤ Δλ0/2 + @test Δλ ≤ Δλ0 / 2 end end diff --git a/test/runtests.jl b/test/runtests.jl index e1e230eb9..4fc054147 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -17,6 +17,7 @@ using Random, StableRNGs using Statistics using StatsBase +using DifferentiationInterface using AdvancedVI const PROGRESS = haskey(ENV, "PROGRESS") From 79cc048af6d30055a72aa202896d2e5ce7263999 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 19 May 2026 18:02:56 +0100 Subject: [PATCH 02/31] Bump AbstractPPL@0.15, add LDP+Hessian on prepared evaluators MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Switch `value_and_gradient` → `value_and_gradient!!` per AbstractPPL 0.15. - In AdvancedVI core, give every `AbstractPPL.Prepared{<:AbstractADType,<:VectorEvaluator}` a `LogDensityProblems` interface (order-1 fallback) so any AD-backed prep acts as a `LogDensityProblem` without backend-specific wiring. - In `AdvancedVIMooncakeExt`, promote Mooncake-prepared evaluators to order-2 and add `logdensity_gradient_and_hessian` via forward-over-reverse Mooncake. - Simplify `DynamicPPLModelLogDensityFunction` to delegate LDP calls to its inner prep; `capabilities` reads off the prep's own capability so the Hessian branch is exposed exactly when the backend supports it. - Bump compat: AbstractPPL 0.15, DynamicPPL 0.42 (with branch pin until release); add Bijectors branch pin in `test/` for the same reason. Co-Authored-By: Claude Opus 4.7 (1M context) --- Project.toml | 5 ++- ext/AdvancedVIDynamicPPLExt.jl | 56 +++++++++++++++------------------- ext/AdvancedVIMooncakeExt.jl | 22 +++++++++++++ src/AdvancedVI.jl | 36 +++++++++++++++++++--- test/Project.toml | 5 +-- 5 files changed, 83 insertions(+), 41 deletions(-) diff --git a/Project.toml b/Project.toml index 886401ebe..8cffa4dcf 100644 --- a/Project.toml +++ b/Project.toml @@ -33,13 +33,13 @@ AdvancedVIDynamicPPLExt = ["DynamicPPL", "Accessors", "Distributions", "LogDensi [compat] ADTypes = "1" -AbstractPPL = "0.14" +AbstractPPL = "0.15" Accessors = "0.1" ChainRulesCore = "1" DiffResults = "1" Distributions = "0.25.111" DocStringExtensions = "0.8, 0.9" -DynamicPPL = "0.40, 0.41" +DynamicPPL = "0.42" Enzyme = "0.13" FillArrays = "1.3" Functors = "0.4, 0.5" @@ -65,5 +65,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" test = ["Pkg", "Test"] [sources] -AbstractPPL = {url = "https://github.com/TuringLang/AbstractPPL.jl", rev = "evaluator-interface"} DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "adproblems-interface"} diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index 08c4bcf4d..6e5f86eef 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -1,22 +1,12 @@ module AdvancedVIDynamicPPLExt using ADTypes: ADTypes -using Accessors using AdvancedVI: AdvancedVI using AbstractPPL: AbstractPPL -using Distributions: Distributions using DynamicPPL: DynamicPPL using LogDensityProblems: LogDensityProblems using Random -function adtype_capabilities(::Type{Nothing}) - return LogDensityProblems.LogDensityOrder{0}() -end - -function adtype_capabilities(::Type{<:ADTypes.AbstractADType}) - return LogDensityProblems.LogDensityOrder{1}() -end - function logdensity_impl( params, model::DynamicPPL.Model, loglikeadj::Real, varinfo::DynamicPPL.AbstractVarInfo ) @@ -35,31 +25,29 @@ function subsample_dynamicpplmodel( return DynamicPPL.Model{Threaded}(model.f, model.args, new_kwargs, model.context) end +# `LogDensityProblems.capabilities` and the gradient/Hessian methods dispatch +# off `Prep`, so the AD backend's `Prepared` type drives the LDP capability. struct DynamicPPLModelLogDensityFunction{ Model<:DynamicPPL.Model, VarInfo<:DynamicPPL.AbstractVarInfo, ADType<:Union{Nothing,ADTypes.AbstractADType}, - PrepGrad, + Prep, } model::Model varinfo::VarInfo adtype::ADType - # Refs are updated in-place by subsample; the prepared AD evaluator reads - # through them on every call, so the prep remains valid across subsampling. model_ref::Ref{Any} loglikeadj_ref::Ref{Float64} - prep_grad::PrepGrad + prep::Prep end function DynamicPPLModelLogDensityFunction( model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo; - use_hessian::Bool=false, adtype::Union{Nothing,ADTypes.AbstractADType}=nothing, loglikeadj::Real=1.0, subsampling::Union{Nothing,AdvancedVI.AbstractSubsampling}=nothing, ) - use_hessian && @warn "`use_hessian` is no longer supported and will be ignored." model_sub = if isnothing(subsampling) model else @@ -69,24 +57,20 @@ function DynamicPPLModelLogDensityFunction( subsample_dynamicpplmodel(model, batch) end - params = [val for val in varinfo[:]] - cap = adtype_capabilities(typeof(adtype)) + params = collect(varinfo[:]) model_ref = Ref{Any}(model_sub) loglikeadj_ref = Ref{Float64}(float(loglikeadj)) - prep_grad = if cap >= LogDensityProblems.LogDensityOrder{1}() - AbstractPPL.prepare( - adtype, - params -> logdensity_impl(params, model_ref[], loglikeadj_ref[], varinfo), - params, - ) - else + prep = if isnothing(adtype) nothing + else + f = params -> logdensity_impl(params, model_ref[], loglikeadj_ref[], varinfo) + AbstractPPL.prepare(adtype, f, params) end return DynamicPPLModelLogDensityFunction( - model, varinfo, adtype, model_ref, loglikeadj_ref, prep_grad + model, varinfo, adtype, model_ref, loglikeadj_ref, prep ) end @@ -97,19 +81,25 @@ end function LogDensityProblems.logdensity_and_gradient( prob::DynamicPPLModelLogDensityFunction, params ) - return AbstractPPL.value_and_gradient(prob.prep_grad, params) + return LogDensityProblems.logdensity_and_gradient(prob.prep, params) +end + +function LogDensityProblems.logdensity_gradient_and_hessian( + prob::DynamicPPLModelLogDensityFunction, params +) + return LogDensityProblems.logdensity_gradient_and_hessian(prob.prep, params) end function LogDensityProblems.capabilities( - ::Type{<:DynamicPPLModelLogDensityFunction{M,V,Nothing,G}} -) where {M,V,G} + ::Type{<:DynamicPPLModelLogDensityFunction{M,V,Nothing,P}} +) where {M,V,P} return LogDensityProblems.LogDensityOrder{0}() end function LogDensityProblems.capabilities( - ::Type{<:DynamicPPLModelLogDensityFunction{M,V,<:ADTypes.AbstractADType,G}} -) where {M,V,G} - return LogDensityProblems.LogDensityOrder{1}() + ::Type{<:DynamicPPLModelLogDensityFunction{M,V,A,P}} +) where {M,V,A<:ADTypes.AbstractADType,P} + return LogDensityProblems.capabilities(P) end function LogDensityProblems.dimension(prob::DynamicPPLModelLogDensityFunction) @@ -132,6 +122,8 @@ function AdvancedVI.subsample(prob::DynamicPPLModelLogDensityFunction, batch) model_sub = subsample_dynamicpplmodel(model, batch) loglikeadj = n_datapoints / batchsize + # Mutates the refs so the previously prepared AD evaluator keeps reading + # the latest batch without needing a re-prepare. prob.model_ref[] = model_sub prob.loglikeadj_ref[] = loglikeadj diff --git a/ext/AdvancedVIMooncakeExt.jl b/ext/AdvancedVIMooncakeExt.jl index 605f77bfa..9b562aad8 100644 --- a/ext/AdvancedVIMooncakeExt.jl +++ b/ext/AdvancedVIMooncakeExt.jl @@ -1,5 +1,8 @@ module AdvancedVIMooncakeExt +using ADTypes: AutoMooncake, AutoMooncakeForward +using AbstractPPL: AbstractPPL +using AbstractPPL.Evaluators: Prepared, VectorEvaluator using AdvancedVI using LogDensityProblems using Mooncake @@ -31,4 +34,23 @@ function Mooncake.rrule!!( return Mooncake.zero_fcodual(ℓπ), logdensity_pb end +const _MooncakePrepared = Prepared{<:AutoMooncake,<:VectorEvaluator} + +# Order-1 LDP methods are inherited from the AbstractADType fallback in +# AdvancedVI core. +function LogDensityProblems.capabilities(::Type{<:_MooncakePrepared}) + LogDensityProblems.LogDensityOrder{2}() +end + +# Mooncake forward-over-reverse Hessian: a fresh forward-mode Jacobian cache +# is built per call, so this is fine for occasional use but costly inside a +# tight per-sample loop. +function LogDensityProblems.logdensity_gradient_and_hessian(p::_MooncakePrepared, x) + val, grad = LogDensityProblems.logdensity_and_gradient(p, x) + grad_fn = y -> LogDensityProblems.logdensity_and_gradient(p, y)[2] + fwd_jac = AbstractPPL.prepare(AutoMooncakeForward(), grad_fn, x) + _, H = AbstractPPL.value_and_jacobian!!(fwd_jac, x) + return val, grad, copy(H) +end + end diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index a0ce225de..82c865a5e 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -18,14 +18,15 @@ using LogDensityProblems using ADTypes using DiffResults using AbstractPPL: AbstractPPL +using AbstractPPL.Evaluators: Prepared, VectorEvaluator using ChainRulesCore: ChainRulesCore using FillArrays using StatsBase -# Holds the AbstractPPL prepared evaluator together with the aux Ref so that -# _value_and_gradient! can update aux before every evaluation. +# `aux` is captured by Ref so the same prepared evaluator can be reused after +# aux changes — re-preparing per call would defeat the cache. struct _VIGradPrep{P,R} prepared::P aux_ref::R @@ -59,7 +60,7 @@ function _value_and_gradient!( f, out::DiffResults.MutableDiffResult, ad::ADTypes.AbstractADType, x, aux ) prepared = AbstractPPL.prepare(ad, Base.Fix2(f, aux), x) - val, grad = AbstractPPL.value_and_gradient(prepared, x) + val, grad = AbstractPPL.value_and_gradient!!(prepared, x) DiffResults.value!(out, val) copyto!(DiffResults.gradient(out), grad) return out @@ -74,7 +75,7 @@ function _value_and_gradient!( aux, ) prep.aux_ref[] = aux - val, grad = AbstractPPL.value_and_gradient(prep.prepared, x) + val, grad = AbstractPPL.value_and_gradient!!(prep.prepared, x) DiffResults.value!(out, val) copyto!(DiffResults.gradient(out), grad) return out @@ -110,6 +111,33 @@ This is an indirection for handling the type stability of `restructure`, as some """ restructure_ad_forward(::ADTypes.AbstractADType, restructure, params) = restructure(params) +# Gradient-only LDP fallback for any AD-prepared evaluator; backend extensions +# override `capabilities` and add `logdensity_gradient_and_hessian` if they can. +function LogDensityProblems.capabilities( + ::Type{<:Prepared{<:ADTypes.AbstractADType,<:VectorEvaluator}} +) + LogDensityProblems.LogDensityOrder{1}() +end + +function LogDensityProblems.dimension( + p::Prepared{<:ADTypes.AbstractADType,<:VectorEvaluator} +) + p.evaluator.dim +end + +function LogDensityProblems.logdensity( + p::Prepared{<:ADTypes.AbstractADType,<:VectorEvaluator}, x +) + p(x) +end + +function LogDensityProblems.logdensity_and_gradient( + p::Prepared{<:ADTypes.AbstractADType,<:VectorEvaluator}, x +) + val, grad = AbstractPPL.value_and_gradient!!(p, x) + return val, copy(grad) +end + include("mixedad_logdensity.jl") # Variational Families diff --git a/test/Project.toml b/test/Project.toml index 6aad881ab..d85b47d4d 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -2,6 +2,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" +Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" @@ -26,15 +27,15 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [sources] -AbstractPPL = {url = "https://github.com/TuringLang/AbstractPPL.jl", rev = "evaluator-interface"} DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "adproblems-interface"} +Bijectors = {url = "https://github.com/TuringLang/Bijectors.jl", rev = "replace-di-with-abstractppl"} [compat] ADTypes = "0.2.1, 1" DiffResults = "1" DifferentiationInterface = "0.6, 0.7" Distributions = "0.25.111" -DynamicPPL = "0.40, 0.41" +DynamicPPL = "0.42" Enzyme = "0.13, 0.14, 0.15" FillArrays = "1.6.1" ForwardDiff = "0.10.36, 1" From 92124d6aa0a117aee86a1988d1fb281cbb849556 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 19 May 2026 18:14:55 +0100 Subject: [PATCH 03/31] Revert formatting-only changes pulled in by the migration `git diff main --stat` was full of `return` keyword and `*` spacing adjustments unrelated to the AbstractPPL 0.15 switch. Restoring those files to `main`'s state shrinks the review surface to just the load-bearing changes (AbstractPPL/DynamicPPL bump, LDP+Hessian wiring, stale-DI comment fixes). Co-Authored-By: Claude Opus 4.7 (1M context) --- src/AdvancedVI.jl | 10 +++++----- src/algorithms/abstractobjective.jl | 2 +- src/algorithms/common.jl | 2 +- src/algorithms/fisherminbatchmatch.jl | 18 +++++++++--------- src/algorithms/klminnaturalgraddescent.jl | 10 +++++----- src/algorithms/klminsqrtnaturalgraddescent.jl | 6 +++--- src/algorithms/klminwassfwdbwd.jl | 6 +++--- src/optimization/rules.jl | 8 ++------ test/integration/dynamicppl.jl | 14 +++++++------- 9 files changed, 36 insertions(+), 40 deletions(-) diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 82c865a5e..5c1105a3b 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -41,9 +41,9 @@ Evaluate the value and gradient of a function `f` at `x` using the automatic dif `f` may receive auxiliary input as `f(x,aux)`. # Arguments -- `ad::ADTypes.AbstractADType`: +- `ad::ADTypes.AbstractADType`: automatic differentiation backend. Currently supports - `ADTypes.AutoZygote()`, `ADTypes.ForwardDiff()`, `ADTypes.ReverseDiff()`, + `ADTypes.AutoZygote()`, `ADTypes.ForwardDiff()`, `ADTypes.ReverseDiff()`, `ADTypes.AutoMooncake()` and `ADTypes.AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), @@ -280,7 +280,7 @@ function step( objargs...; kwargs..., ) - return nothing + nothing end """ @@ -318,11 +318,11 @@ Please refer to the respective documentation of each algorithm for more info. function estimate_objective( ::Random.AbstractRNG, ::AbstractVariationalAlgorithm, q, prob; kwargs... ) - return nothing + nothing end function estimate_objective(alg::AbstractVariationalAlgorithm, q, prob; kwargs...) - return estimate_objective(Random.default_rng(), alg, q, prob; kwargs...) + estimate_objective(Random.default_rng(), alg, q, prob; kwargs...) end export estimate_objective diff --git a/src/algorithms/abstractobjective.jl b/src/algorithms/abstractobjective.jl index ff027bb54..65316c7e7 100644 --- a/src/algorithms/abstractobjective.jl +++ b/src/algorithms/abstractobjective.jl @@ -31,7 +31,7 @@ function init( ::Any, ::Any, ) - return nothing + nothing end """ diff --git a/src/algorithms/common.jl b/src/algorithms/common.jl index 96187fe6a..0b99ff0d0 100644 --- a/src/algorithms/common.jl +++ b/src/algorithms/common.jl @@ -116,5 +116,5 @@ function step( ) info = !isnothing(info′) ? merge(info′, info) : info end - return state, false, info + state, false, info end diff --git a/src/algorithms/fisherminbatchmatch.jl b/src/algorithms/fisherminbatchmatch.jl index f3ea6fd8d..b794a12af 100644 --- a/src/algorithms/fisherminbatchmatch.jl +++ b/src/algorithms/fisherminbatchmatch.jl @@ -89,14 +89,14 @@ function rand_batch_match_samples_with_objective!( μ = q.location C = q.scale u = Random.randn!(rng, u_buf) - z = C * u .+ μ + z = C*u .+ μ logπ_sum = zero(eltype(μ)) for b in 1:n_samples logπb, gb = LogDensityProblems.logdensity_and_gradient(prob, view(z, :, b)) grad_buf[:, b] = gb logπ_sum += logπb end - logπ_avg = logπ_sum / n_samples + logπ_avg = logπ_sum/n_samples # Estimate objective values # @@ -105,7 +105,7 @@ function rand_batch_match_samples_with_objective!( # = E[| C' ( -(CC')\((Cu + μ) - μ) - ∇logπ(z)) |^2] (z = Cu + μ) # = E[| C' ( -(CC')\(Cu) - ∇logπ(z)) |^2] # = E[| -u - C'∇logπ(z)) |^2] - fisher = sum(abs2, -u_buf - (C' * grad_buf)) / n_samples + fisher = sum(abs2, -u_buf - (C'*grad_buf))/n_samples return u_buf, z, grad_buf, fisher, logπ_avg end @@ -145,13 +145,13 @@ function step( gbar, Γ = mean_and_cov(grad_buf, 2) μmz = μ - zbar - λ = convert(eltype(μ), d * n_samples / iteration) + λ = convert(eltype(μ), d*n_samples / iteration) - U = Symmetric(λ * Γ + (λ / (1 + λ) * gbar) * gbar') - V = Symmetric(Σ + λ * C + (λ / (1 + λ) * μmz) * μmz') + U = Symmetric(λ*Γ + (λ/(1 + λ)*gbar)*gbar') + V = Symmetric(Σ + λ*C + (λ/(1 + λ)*μmz)*μmz') - Σ′ = Hermitian(2 * V / (I + real(sqrt(I + 4 * U * V)))) - μ′ = 1 / (1 + λ) * μ + λ / (1 + λ) * (Σ′ * gbar + zbar) + Σ′ = Hermitian(2*V/(I + real(sqrt(I + 4*U*V)))) + μ′ = 1/(1 + λ)*μ + λ/(1 + λ)*(Σ′*gbar + zbar) q′ = MvLocationScale(μ′[:, 1], cholesky(Σ′).L, q.dist) elbo = logπ_avg + entropy(q) @@ -163,7 +163,7 @@ function step( info′ = callback(; rng, iteration, q, state) info = !isnothing(info′) ? merge(info′, info) : info end - return state, false, info + state, false, info end """ diff --git a/src/algorithms/klminnaturalgraddescent.jl b/src/algorithms/klminnaturalgraddescent.jl index 596cd372c..c101537a4 100644 --- a/src/algorithms/klminnaturalgraddescent.jl +++ b/src/algorithms/klminnaturalgraddescent.jl @@ -81,10 +81,10 @@ function init( grad_buf = Vector{eltype(q_init.location)}(undef, n_dims) hess_buf = Matrix{eltype(q_init.location)}(undef, n_dims, n_dims) scale = q_init.scale - qcov = Hermitian(scale * scale') + qcov = Hermitian(scale*scale') scale_inv = inv(scale) prec_chol = scale_inv' - prec = Hermitian(prec_chol * prec_chol') + prec = Hermitian(prec_chol*prec_chol') return KLMinNaturalGradDescentState( q_init, prob, prec, qcov, 0, sub_st, grad_buf, hess_buf ) @@ -127,7 +127,7 @@ function step( # Handling the positive-definite constraint in the Bayesian learning rule. # In ICML 2020. G_hat = S - (-hess_buf) - Hermitian(S - η * G_hat + η^2 / 2 * G_hat * qcov * G_hat) + Hermitian(S - η*G_hat + η^2/2*G_hat*qcov*G_hat) else Hermitian(((1 - η) * S + η * (-hess_buf))) end @@ -136,7 +136,7 @@ function step( prec_chol = cholesky(S′).L prec_chol_inv = inv(prec_chol) scale = prec_chol_inv' - qcov = Hermitian(scale * scale') + qcov = Hermitian(scale*scale') q′ = MvLocationScale(m′, scale, q.dist) state = KLMinNaturalGradDescentState( @@ -149,7 +149,7 @@ function step( info′ = callback(; rng, iteration, q=q′, info) info = !isnothing(info′) ? merge(info′, info) : info end - return state, false, info + state, false, info end """ diff --git a/src/algorithms/klminsqrtnaturalgraddescent.jl b/src/algorithms/klminsqrtnaturalgraddescent.jl index 3b29c3b2b..622c111c3 100644 --- a/src/algorithms/klminsqrtnaturalgraddescent.jl +++ b/src/algorithms/klminsqrtnaturalgraddescent.jl @@ -105,8 +105,8 @@ function step( rng, q, n_samples, grad_buf, hess_buf, prob_sub ) - CtHCmI = C' * (-hess_buf) * C - I - CtHCmI_tril = LowerTriangular(tril(CtHCmI) - Diagonal(diag(CtHCmI)) / 2) + CtHCmI = C'*(-hess_buf)*C - I + CtHCmI_tril = LowerTriangular(tril(CtHCmI) - Diagonal(diag(CtHCmI))/2) m′ = m - η * C * (C' * -grad_buf) C′ = C - η * C * CtHCmI_tril @@ -123,7 +123,7 @@ function step( info′ = callback(; rng, iteration, q=q′, info) info = !isnothing(info′) ? merge(info′, info) : info end - return state, false, info + state, false, info end """ diff --git a/src/algorithms/klminwassfwdbwd.jl b/src/algorithms/klminwassfwdbwd.jl index 40577bfc4..602f4be41 100644 --- a/src/algorithms/klminwassfwdbwd.jl +++ b/src/algorithms/klminwassfwdbwd.jl @@ -104,10 +104,10 @@ function step( m′ = m - η * (-grad_buf) M = I - η * (-hess_buf') - Σ_half = Hermitian(M * Σ * M') + Σ_half = Hermitian(M*Σ*M') # Compute the JKO proximal operator - Σ′ = (Σ_half + 2 * η * I + sqrt(Hermitian(Σ_half * (Σ_half + 4 * η * I)))) / 2 + Σ′ = (Σ_half + 2*η*I + sqrt(Hermitian(Σ_half*(Σ_half + 4*η*I))))/2 q′ = MvLocationScale(m′, cholesky(Σ′).L, q.dist) state = KLMinWassFwdBwdState(q′, prob, Σ′, iteration, sub_st′, grad_buf, hess_buf) @@ -118,7 +118,7 @@ function step( info′ = callback(; rng, iteration, q=q′, info) info = !isnothing(info′) ? merge(info′, info) : info end - return state, false, info + state, false, info end """ diff --git a/src/optimization/rules.jl b/src/optimization/rules.jl index 5015f4ebe..2025632e3 100644 --- a/src/optimization/rules.jl +++ b/src/optimization/rules.jl @@ -18,9 +18,7 @@ Optimisers.@def struct DoWG <: Optimisers.AbstractRule alpha = 1e-6 end -function Optimisers.init(o::DoWG, x::AbstractArray{T}) where {T} - return (copy(x), zero(T), T(o.alpha) * (1 + norm(x))) -end +Optimisers.init(o::DoWG, x::AbstractArray{T}) where {T} = (copy(x), zero(T), T(o.alpha)*(1 + norm(x))) function Optimisers.apply!(::DoWG, state, x::AbstractArray{T}, dx) where {T} x0, v, r = state @@ -49,9 +47,7 @@ Optimisers.@def struct DoG <: Optimisers.AbstractRule alpha = 1e-6 end -function Optimisers.init(o::DoG, x::AbstractArray{T}) where {T} - return (copy(x), zero(T), T(o.alpha) * (1 + norm(x))) -end +Optimisers.init(o::DoG, x::AbstractArray{T}) where {T} = (copy(x), zero(T), T(o.alpha)*(1 + norm(x))) function Optimisers.apply!(::DoG, state, x::AbstractArray{T}, dx) where {T} x0, v, r = state diff --git a/test/integration/dynamicppl.jl b/test/integration/dynamicppl.jl index a3ed01d20..43a38d3b8 100644 --- a/test/integration/dynamicppl.jl +++ b/test/integration/dynamicppl.jl @@ -1,7 +1,7 @@ @testset "DynamicPPL" begin DynamicPPL.@model function normal(μ) - return x ~ MvNormal(μ, I) + x ~ MvNormal(μ, I) end DynamicPPL.@model function normal_subsampled(μs; datapoints=1:size(μs, 2)) @@ -22,18 +22,18 @@ alg = KLMinRepGradProxDescent(AD) d = LogDensityProblems.dimension(prob) - q0 = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6 * I, d, d))) + q0 = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6*I, d, d))) q, _, _ = AdvancedVI.optimize(alg, 1000, prob, q0; show_progress=false) Δλ0 = sum(abs2, q0.location - μ_true) Δλ = sum(abs2, q.location - μ_true) - @test Δλ ≤ Δλ0 / 2 + @test Δλ ≤ Δλ0/2 end @testset "subsampling" begin n_data = 32 - μs = 3 * randn(2, n_data) - μ_true = mean(μs; dims=2)[:, 1] + μs = 3*randn(2, n_data) + μ_true = mean(μs, dims=2)[:, 1] model = normal_subsampled(μs) vi = DynamicPPL.VarInfo(model) @@ -48,11 +48,11 @@ alg = KLMinRepGradProxDescent(AD; subsampling) d = LogDensityProblems.dimension(prob) - q0 = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6 * I, d, d))) + q0 = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6*I, d, d))) q, _, _ = AdvancedVI.optimize(alg, 1000, prob, q0; show_progress=false) Δλ0 = sum(abs2, q0.location - μ_true) Δλ = sum(abs2, q.location - μ_true) - @test Δλ ≤ Δλ0 / 2 + @test Δλ ≤ Δλ0/2 end end From 6da323210bf150e640ad52c050264a380307c2cd Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 19 May 2026 20:54:34 +0100 Subject: [PATCH 04/31] Move LDP+Hessian-on-Prepared work to a follow-up branch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keeps this branch focused on the AbstractPPL evaluator-interface migration: the LDP fallback methods on `Prepared`, the Mooncake forward-over-reverse Hessian, and the DynamicPPLModelLogDensityFunction LDP-delegation refactor have been moved to the `ldp-on-prepared-hessian` branch (branched from `main` and carrying both the migration and the feature work). What stays here: - Switch DI → AbstractPPL.prepare / value_and_gradient!! in core. - DynamicPPLModelLogDensityFunction uses AbstractPPL.prepare instead of DI.prepare_gradient (Hessian path is dropped as in the original migration commit; `use_hessian=true` warns and is ignored). - Compat bumps: AbstractPPL@0.15, DynamicPPL@0.42, plus Bijectors source pin in `test/`. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AdvancedVIDynamicPPLExt.jl | 56 +++++++++++++++++++--------------- ext/AdvancedVIMooncakeExt.jl | 22 ------------- src/AdvancedVI.jl | 28 ----------------- 3 files changed, 32 insertions(+), 74 deletions(-) diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index 6e5f86eef..af3839bc6 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -1,12 +1,22 @@ module AdvancedVIDynamicPPLExt using ADTypes: ADTypes +using Accessors using AdvancedVI: AdvancedVI using AbstractPPL: AbstractPPL +using Distributions: Distributions using DynamicPPL: DynamicPPL using LogDensityProblems: LogDensityProblems using Random +function adtype_capabilities(::Type{Nothing}) + return LogDensityProblems.LogDensityOrder{0}() +end + +function adtype_capabilities(::Type{<:ADTypes.AbstractADType}) + return LogDensityProblems.LogDensityOrder{1}() +end + function logdensity_impl( params, model::DynamicPPL.Model, loglikeadj::Real, varinfo::DynamicPPL.AbstractVarInfo ) @@ -25,29 +35,31 @@ function subsample_dynamicpplmodel( return DynamicPPL.Model{Threaded}(model.f, model.args, new_kwargs, model.context) end -# `LogDensityProblems.capabilities` and the gradient/Hessian methods dispatch -# off `Prep`, so the AD backend's `Prepared` type drives the LDP capability. struct DynamicPPLModelLogDensityFunction{ Model<:DynamicPPL.Model, VarInfo<:DynamicPPL.AbstractVarInfo, ADType<:Union{Nothing,ADTypes.AbstractADType}, - Prep, + PrepGrad, } model::Model varinfo::VarInfo adtype::ADType + # Refs are updated in-place by subsample; the prepared AD evaluator reads + # through them on every call, so the prep remains valid across subsampling. model_ref::Ref{Any} loglikeadj_ref::Ref{Float64} - prep::Prep + prep_grad::PrepGrad end function DynamicPPLModelLogDensityFunction( model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo; + use_hessian::Bool=false, adtype::Union{Nothing,ADTypes.AbstractADType}=nothing, loglikeadj::Real=1.0, subsampling::Union{Nothing,AdvancedVI.AbstractSubsampling}=nothing, ) + use_hessian && @warn "`use_hessian` is no longer supported and will be ignored." model_sub = if isnothing(subsampling) model else @@ -57,20 +69,24 @@ function DynamicPPLModelLogDensityFunction( subsample_dynamicpplmodel(model, batch) end - params = collect(varinfo[:]) + params = [val for val in varinfo[:]] + cap = adtype_capabilities(typeof(adtype)) model_ref = Ref{Any}(model_sub) loglikeadj_ref = Ref{Float64}(float(loglikeadj)) - prep = if isnothing(adtype) - nothing + prep_grad = if cap >= LogDensityProblems.LogDensityOrder{1}() + AbstractPPL.prepare( + adtype, + params -> logdensity_impl(params, model_ref[], loglikeadj_ref[], varinfo), + params, + ) else - f = params -> logdensity_impl(params, model_ref[], loglikeadj_ref[], varinfo) - AbstractPPL.prepare(adtype, f, params) + nothing end return DynamicPPLModelLogDensityFunction( - model, varinfo, adtype, model_ref, loglikeadj_ref, prep + model, varinfo, adtype, model_ref, loglikeadj_ref, prep_grad ) end @@ -81,25 +97,19 @@ end function LogDensityProblems.logdensity_and_gradient( prob::DynamicPPLModelLogDensityFunction, params ) - return LogDensityProblems.logdensity_and_gradient(prob.prep, params) -end - -function LogDensityProblems.logdensity_gradient_and_hessian( - prob::DynamicPPLModelLogDensityFunction, params -) - return LogDensityProblems.logdensity_gradient_and_hessian(prob.prep, params) + return AbstractPPL.value_and_gradient!!(prob.prep_grad, params) end function LogDensityProblems.capabilities( - ::Type{<:DynamicPPLModelLogDensityFunction{M,V,Nothing,P}} -) where {M,V,P} + ::Type{<:DynamicPPLModelLogDensityFunction{M,V,Nothing,G}} +) where {M,V,G} return LogDensityProblems.LogDensityOrder{0}() end function LogDensityProblems.capabilities( - ::Type{<:DynamicPPLModelLogDensityFunction{M,V,A,P}} -) where {M,V,A<:ADTypes.AbstractADType,P} - return LogDensityProblems.capabilities(P) + ::Type{<:DynamicPPLModelLogDensityFunction{M,V,<:ADTypes.AbstractADType,G}} +) where {M,V,G} + return LogDensityProblems.LogDensityOrder{1}() end function LogDensityProblems.dimension(prob::DynamicPPLModelLogDensityFunction) @@ -122,8 +132,6 @@ function AdvancedVI.subsample(prob::DynamicPPLModelLogDensityFunction, batch) model_sub = subsample_dynamicpplmodel(model, batch) loglikeadj = n_datapoints / batchsize - # Mutates the refs so the previously prepared AD evaluator keeps reading - # the latest batch without needing a re-prepare. prob.model_ref[] = model_sub prob.loglikeadj_ref[] = loglikeadj diff --git a/ext/AdvancedVIMooncakeExt.jl b/ext/AdvancedVIMooncakeExt.jl index 9b562aad8..605f77bfa 100644 --- a/ext/AdvancedVIMooncakeExt.jl +++ b/ext/AdvancedVIMooncakeExt.jl @@ -1,8 +1,5 @@ module AdvancedVIMooncakeExt -using ADTypes: AutoMooncake, AutoMooncakeForward -using AbstractPPL: AbstractPPL -using AbstractPPL.Evaluators: Prepared, VectorEvaluator using AdvancedVI using LogDensityProblems using Mooncake @@ -34,23 +31,4 @@ function Mooncake.rrule!!( return Mooncake.zero_fcodual(ℓπ), logdensity_pb end -const _MooncakePrepared = Prepared{<:AutoMooncake,<:VectorEvaluator} - -# Order-1 LDP methods are inherited from the AbstractADType fallback in -# AdvancedVI core. -function LogDensityProblems.capabilities(::Type{<:_MooncakePrepared}) - LogDensityProblems.LogDensityOrder{2}() -end - -# Mooncake forward-over-reverse Hessian: a fresh forward-mode Jacobian cache -# is built per call, so this is fine for occasional use but costly inside a -# tight per-sample loop. -function LogDensityProblems.logdensity_gradient_and_hessian(p::_MooncakePrepared, x) - val, grad = LogDensityProblems.logdensity_and_gradient(p, x) - grad_fn = y -> LogDensityProblems.logdensity_and_gradient(p, y)[2] - fwd_jac = AbstractPPL.prepare(AutoMooncakeForward(), grad_fn, x) - _, H = AbstractPPL.value_and_jacobian!!(fwd_jac, x) - return val, grad, copy(H) -end - end diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 5c1105a3b..765ef2977 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -18,7 +18,6 @@ using LogDensityProblems using ADTypes using DiffResults using AbstractPPL: AbstractPPL -using AbstractPPL.Evaluators: Prepared, VectorEvaluator using ChainRulesCore: ChainRulesCore using FillArrays @@ -111,33 +110,6 @@ This is an indirection for handling the type stability of `restructure`, as some """ restructure_ad_forward(::ADTypes.AbstractADType, restructure, params) = restructure(params) -# Gradient-only LDP fallback for any AD-prepared evaluator; backend extensions -# override `capabilities` and add `logdensity_gradient_and_hessian` if they can. -function LogDensityProblems.capabilities( - ::Type{<:Prepared{<:ADTypes.AbstractADType,<:VectorEvaluator}} -) - LogDensityProblems.LogDensityOrder{1}() -end - -function LogDensityProblems.dimension( - p::Prepared{<:ADTypes.AbstractADType,<:VectorEvaluator} -) - p.evaluator.dim -end - -function LogDensityProblems.logdensity( - p::Prepared{<:ADTypes.AbstractADType,<:VectorEvaluator}, x -) - p(x) -end - -function LogDensityProblems.logdensity_and_gradient( - p::Prepared{<:ADTypes.AbstractADType,<:VectorEvaluator}, x -) - val, grad = AbstractPPL.value_and_gradient!!(p, x) - return val, copy(grad) -end - include("mixedad_logdensity.jl") # Variational Families From d106e7b373c7a6748ff4ab4e19b7a0a8e4d9acdf Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 19 May 2026 22:14:21 +0100 Subject: [PATCH 05/31] Restore Hessian support via AbstractPPL hg/hessian-order MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pins AbstractPPL to `hg/hessian-order`, which adds `prepare(adtype, f, x; order=2)` and `value_gradient_and_hessian!!`. `DynamicPPLModelLogDensityFunction` goes back to main's `prep_grad` + `prep_hess` shape (default `use_hessian=true`), so the diff reads as a clean DI → AbstractPPL swap rather than a redesign. `use_hessian` falls back to gradient-only when the AD backend refuses `order=2` (MethodError only; other errors propagate). Co-Authored-By: Claude Opus 4.7 (1M context) --- Project.toml | 1 + ext/AdvancedVIDynamicPPLExt.jl | 82 +++++++++++++++++++++------------- src/AdvancedVI.jl | 4 +- test/Project.toml | 1 + 4 files changed, 54 insertions(+), 34 deletions(-) diff --git a/Project.toml b/Project.toml index 8cffa4dcf..264f66372 100644 --- a/Project.toml +++ b/Project.toml @@ -65,4 +65,5 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" test = ["Pkg", "Test"] [sources] +AbstractPPL = {url = "https://github.com/TuringLang/AbstractPPL.jl", rev = "hg/hessian-order"} DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "adproblems-interface"} diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index af3839bc6..95a28ba3d 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -1,17 +1,13 @@ module AdvancedVIDynamicPPLExt using ADTypes: ADTypes -using Accessors using AdvancedVI: AdvancedVI using AbstractPPL: AbstractPPL -using Distributions: Distributions using DynamicPPL: DynamicPPL using LogDensityProblems: LogDensityProblems using Random -function adtype_capabilities(::Type{Nothing}) - return LogDensityProblems.LogDensityOrder{0}() -end +adtype_capabilities(::Type{Nothing}) = LogDensityProblems.LogDensityOrder{0}() function adtype_capabilities(::Type{<:ADTypes.AbstractADType}) return LogDensityProblems.LogDensityOrder{1}() @@ -35,31 +31,33 @@ function subsample_dynamicpplmodel( return DynamicPPL.Model{Threaded}(model.f, model.args, new_kwargs, model.context) end +# `model_ref`/`loglikeadj_ref` are mutated in place by `subsample`; the closure +# inside `prep_grad`/`prep_hess` reads through them so the prep stays valid +# across subsampling steps (AbstractPPL bakes the closure into the prep, unlike +# DI's `Constant` which can be rebound at call time). struct DynamicPPLModelLogDensityFunction{ Model<:DynamicPPL.Model, VarInfo<:DynamicPPL.AbstractVarInfo, ADType<:Union{Nothing,ADTypes.AbstractADType}, PrepGrad, + PrepHess, } - model::Model - varinfo::VarInfo - adtype::ADType - # Refs are updated in-place by subsample; the prepared AD evaluator reads - # through them on every call, so the prep remains valid across subsampling. model_ref::Ref{Any} loglikeadj_ref::Ref{Float64} + varinfo::VarInfo + adtype::ADType prep_grad::PrepGrad + prep_hess::PrepHess end function DynamicPPLModelLogDensityFunction( model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo; - use_hessian::Bool=false, + use_hessian::Bool=true, adtype::Union{Nothing,ADTypes.AbstractADType}=nothing, loglikeadj::Real=1.0, subsampling::Union{Nothing,AdvancedVI.AbstractSubsampling}=nothing, ) - use_hessian && @warn "`use_hessian` is no longer supported and will be ignored." model_sub = if isnothing(subsampling) model else @@ -69,24 +67,36 @@ function DynamicPPLModelLogDensityFunction( subsample_dynamicpplmodel(model, batch) end - params = [val for val in varinfo[:]] - cap = adtype_capabilities(typeof(adtype)) - + params = collect(varinfo[:]) model_ref = Ref{Any}(model_sub) loglikeadj_ref = Ref{Float64}(float(loglikeadj)) + f = params -> logdensity_impl(params, model_ref[], loglikeadj_ref[], varinfo) + cap = adtype_capabilities(typeof(adtype)) prep_grad = if cap >= LogDensityProblems.LogDensityOrder{1}() - AbstractPPL.prepare( - adtype, - params -> logdensity_impl(params, model_ref[], loglikeadj_ref[], varinfo), - params, - ) + AbstractPPL.prepare(adtype, f, params) else nothing end - - return DynamicPPLModelLogDensityFunction( - model, varinfo, adtype, model_ref, loglikeadj_ref, prep_grad + prep_hess = if cap >= LogDensityProblems.LogDensityOrder{1}() && use_hessian + try + AbstractPPL.prepare(adtype, f, params; order=2) + catch err + err isa MethodError || rethrow() + @warn "The selected AD backend does not support `AbstractPPL.prepare(...; order=2)`. AdvancedVI will treat the model as first-order only." + nothing + end + else + nothing + end + return DynamicPPLModelLogDensityFunction{ + typeof(model), + typeof(varinfo), + typeof(adtype), + typeof(prep_grad), + typeof(prep_hess), + }( + model_ref, loglikeadj_ref, varinfo, adtype, prep_grad, prep_hess ) end @@ -97,19 +107,27 @@ end function LogDensityProblems.logdensity_and_gradient( prob::DynamicPPLModelLogDensityFunction, params ) - return AbstractPPL.value_and_gradient!!(prob.prep_grad, params) + val, grad = AbstractPPL.value_and_gradient!!(prob.prep_grad, params) + return val, copy(grad) end -function LogDensityProblems.capabilities( - ::Type{<:DynamicPPLModelLogDensityFunction{M,V,Nothing,G}} -) where {M,V,G} - return LogDensityProblems.LogDensityOrder{0}() +function LogDensityProblems.logdensity_gradient_and_hessian( + prob::DynamicPPLModelLogDensityFunction, params +) + val, grad, H = AbstractPPL.value_gradient_and_hessian!!(prob.prep_hess, params) + return val, copy(grad), copy(H) end function LogDensityProblems.capabilities( - ::Type{<:DynamicPPLModelLogDensityFunction{M,V,<:ADTypes.AbstractADType,G}} -) where {M,V,G} - return LogDensityProblems.LogDensityOrder{1}() + ::Type{<:DynamicPPLModelLogDensityFunction{M,V,ADType,PG,PH}} +) where {M,V,ADType<:ADTypes.AbstractADType,PG,PH} + return if PH != Nothing + LogDensityProblems.LogDensityOrder{2}() + elseif PG != Nothing + LogDensityProblems.LogDensityOrder{1}() + else + LogDensityProblems.LogDensityOrder{0}() + end end function LogDensityProblems.dimension(prob::DynamicPPLModelLogDensityFunction) @@ -117,7 +135,7 @@ function LogDensityProblems.dimension(prob::DynamicPPLModelLogDensityFunction) end function AdvancedVI.subsample(prob::DynamicPPLModelLogDensityFunction, batch) - model = prob.model + model = prob.model_ref[] if !haskey(model.defaults, :datapoints) throw( diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 765ef2977..9fd6b64cd 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -24,8 +24,8 @@ using FillArrays using StatsBase -# `aux` is captured by Ref so the same prepared evaluator can be reused after -# aux changes — re-preparing per call would defeat the cache. +# `AbstractPPL.prepare` bakes the closure at prep time, so `aux` is captured +# via a `Ref` that callers mutate before each evaluation. struct _VIGradPrep{P,R} prepared::P aux_ref::R diff --git a/test/Project.toml b/test/Project.toml index d85b47d4d..ae013611c 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -27,6 +27,7 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [sources] +AbstractPPL = {url = "https://github.com/TuringLang/AbstractPPL.jl", rev = "hg/hessian-order"} DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "adproblems-interface"} Bijectors = {url = "https://github.com/TuringLang/Bijectors.jl", rev = "replace-di-with-abstractppl"} From 79db2fb87b90e7e76d3c7ffa737e715cfec4d0b4 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 19 May 2026 22:22:03 +0100 Subject: [PATCH 06/31] Drop _VIGradPrep; thread aux through AbstractPPL context Wrapping the prepared evaluator and the aux `Ref` in a struct was just working around the lack of mutable state on `AbstractPPL.Prepared`. AbstractPPL 0.15's `prepare(...; context=tuple)` keeps the tuple on the evaluator and threads it through every call, so passing `Ref(aux)` in the context stores the ref inside the prep itself; callers mutate `prep.evaluator.context[1][]` before each evaluation. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AdvancedVIDynamicPPLExt.jl | 6 +----- src/AdvancedVI.jl | 22 ++++------------------ 2 files changed, 5 insertions(+), 23 deletions(-) diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index 95a28ba3d..53ed27c35 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -90,11 +90,7 @@ function DynamicPPLModelLogDensityFunction( nothing end return DynamicPPLModelLogDensityFunction{ - typeof(model), - typeof(varinfo), - typeof(adtype), - typeof(prep_grad), - typeof(prep_hess), + typeof(model),typeof(varinfo),typeof(adtype),typeof(prep_grad),typeof(prep_hess) }( model_ref, loglikeadj_ref, varinfo, adtype, prep_grad, prep_hess ) diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 9fd6b64cd..3e32f7359 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -24,13 +24,6 @@ using FillArrays using StatsBase -# `AbstractPPL.prepare` bakes the closure at prep time, so `aux` is captured -# via a `Ref` that callers mutate before each evaluation. -struct _VIGradPrep{P,R} - prepared::P - aux_ref::R -end - # Derivatives """ _value_and_gradient!(f, out, ad, x, aux) @@ -66,15 +59,10 @@ function _value_and_gradient!( end function _value_and_gradient!( - f, - out::DiffResults.MutableDiffResult, - prep::_VIGradPrep, - ad::ADTypes.AbstractADType, - x, - aux, + f, out::DiffResults.MutableDiffResult, prep, ad::ADTypes.AbstractADType, x, aux ) - prep.aux_ref[] = aux - val, grad = AbstractPPL.value_and_gradient!!(prep.prepared, x) + prep.evaluator.context[1][] = aux + val, grad = AbstractPPL.value_and_gradient!!(prep, x) DiffResults.value!(out, val) copyto!(DiffResults.gradient(out), grad) return out @@ -92,9 +80,7 @@ Prepare AD backend for taking gradients of a function `f` at `x` using the autom - `aux`: Auxiliary input passed to `f`. """ function _prepare_gradient(f, ad::ADTypes.AbstractADType, x, aux) - aux_ref = Ref(aux) - prepared = AbstractPPL.prepare(ad, x -> f(x, aux_ref[]), x) - return _VIGradPrep(prepared, aux_ref) + return AbstractPPL.prepare(ad, (x, aref) -> f(x, aref[]), x; context=(Ref(aux),)) end """ From 451e7b706fd2b7749b751213670f787306869293 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 20 May 2026 11:13:12 +0100 Subject: [PATCH 07/31] Parameterise loglikeadj eltype; expand model_ref WHY comment `loglikeadj_ref` now tracks the input numeric type via a `LogLikeAdj<:Real` struct parameter, restoring `Float32`/`BigFloat`/dual support that the hard-pinned `Ref{Float64}` quietly dropped. `subsample` computes the adjustment in the field's eltype to preserve precision across minibatches. The block comment now also flags why `model_ref::Ref{Any}` cannot be tightened (subsample-widened `defaults` NamedTuple), so a future reader doesn't try. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AdvancedVIDynamicPPLExt.jl | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index 53ed27c35..699327fc6 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -35,15 +35,21 @@ end # inside `prep_grad`/`prep_hess` reads through them so the prep stays valid # across subsampling steps (AbstractPPL bakes the closure into the prep, unlike # DI's `Constant` which can be rebound at call time). +# +# `model_ref` is typed `Ref{Any}` because `subsample_dynamicpplmodel` returns a +# `DynamicPPL.Model` whose `defaults` NamedTuple type varies with the batch — a +# typed `Ref{<:DynamicPPL.Model}` would throw on reassignment. The tradeoff is a +# dynamic dispatch on each `prob.model_ref[]` read; do not "tighten" it. struct DynamicPPLModelLogDensityFunction{ Model<:DynamicPPL.Model, + LogLikeAdj<:Real, VarInfo<:DynamicPPL.AbstractVarInfo, ADType<:Union{Nothing,ADTypes.AbstractADType}, PrepGrad, PrepHess, } model_ref::Ref{Any} - loglikeadj_ref::Ref{Float64} + loglikeadj_ref::Ref{LogLikeAdj} varinfo::VarInfo adtype::ADType prep_grad::PrepGrad @@ -69,7 +75,8 @@ function DynamicPPLModelLogDensityFunction( params = collect(varinfo[:]) model_ref = Ref{Any}(model_sub) - loglikeadj_ref = Ref{Float64}(float(loglikeadj)) + adj0 = float(loglikeadj) + loglikeadj_ref = Ref(adj0) f = params -> logdensity_impl(params, model_ref[], loglikeadj_ref[], varinfo) cap = adtype_capabilities(typeof(adtype)) @@ -90,7 +97,12 @@ function DynamicPPLModelLogDensityFunction( nothing end return DynamicPPLModelLogDensityFunction{ - typeof(model),typeof(varinfo),typeof(adtype),typeof(prep_grad),typeof(prep_hess) + typeof(model), + typeof(adj0), + typeof(varinfo), + typeof(adtype), + typeof(prep_grad), + typeof(prep_hess), }( model_ref, loglikeadj_ref, varinfo, adtype, prep_grad, prep_hess ) @@ -115,8 +127,8 @@ function LogDensityProblems.logdensity_gradient_and_hessian( end function LogDensityProblems.capabilities( - ::Type{<:DynamicPPLModelLogDensityFunction{M,V,ADType,PG,PH}} -) where {M,V,ADType<:ADTypes.AbstractADType,PG,PH} + ::Type{<:DynamicPPLModelLogDensityFunction{M,L,V,ADType,PG,PH}} +) where {M,L,V,ADType<:ADTypes.AbstractADType,PG,PH} return if PH != Nothing LogDensityProblems.LogDensityOrder{2}() elseif PG != Nothing @@ -144,7 +156,8 @@ function AdvancedVI.subsample(prob::DynamicPPLModelLogDensityFunction, batch) n_datapoints = length(model.defaults.datapoints) batchsize = length(batch) model_sub = subsample_dynamicpplmodel(model, batch) - loglikeadj = n_datapoints / batchsize + T = eltype(prob.loglikeadj_ref) + loglikeadj = T(n_datapoints) / T(batchsize) prob.model_ref[] = model_sub prob.loglikeadj_ref[] = loglikeadj From 5790cd422612cdc891b14dae828e3092a8d822d2 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 21 May 2026 20:48:50 +0100 Subject: [PATCH 08/31] Drop AbstractPPL and Bijectors source pins; restrict Bijectors compat AbstractPPL 0.15.1 and Bijectors 0.16.0 are now released. The DynamicPPL `adproblems-interface` branch has been updated to require Bijectors 0.16, so keeping the DynamicPPL source pin no longer pulls in an older Bijectors. Restrict the package and test envs to Bijectors 0.16; relax docs to "0.15, 0.16" since NormalizingFlows 0.2.2 still requires 0.15.x. Co-Authored-By: Claude Opus 4.7 (1M context) --- Project.toml | 1 - docs/Project.toml | 2 +- test/Project.toml | 3 +-- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/Project.toml b/Project.toml index 264f66372..8cffa4dcf 100644 --- a/Project.toml +++ b/Project.toml @@ -65,5 +65,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" test = ["Pkg", "Test"] [sources] -AbstractPPL = {url = "https://github.com/TuringLang/AbstractPPL.jl", rev = "hg/hessian-order"} DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "adproblems-interface"} diff --git a/docs/Project.toml b/docs/Project.toml index c05bf9e27..9374d4c52 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -26,7 +26,7 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c" ADTypes = "1" Accessors = "0.1" AdvancedVI = "0.7, 0.6" -Bijectors = "0.13.6, 0.14, 0.15" +Bijectors = "0.15, 0.16" DataFrames = "1" DifferentiationInterface = "0.7" Distributions = "0.25" diff --git a/test/Project.toml b/test/Project.toml index d7583ac7b..6114a8fdf 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -26,12 +26,11 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" [sources] -AbstractPPL = {url = "https://github.com/TuringLang/AbstractPPL.jl", rev = "hg/hessian-order"} DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "adproblems-interface"} -Bijectors = {url = "https://github.com/TuringLang/Bijectors.jl", rev = "replace-di-with-abstractppl"} [compat] ADTypes = "0.2.1, 1" +Bijectors = "0.16" DiffResults = "1" DifferentiationInterface = "0.6, 0.7" Distributions = "0.25.111" From 6882db840e6e1e2e691ea51bb90a6b60a8b9b206 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 20:31:33 +0100 Subject: [PATCH 09/31] Repoint DynamicPPL source pin to main The adproblems-interface branch was merged into DynamicPPL main (#1363) and deleted, breaking CI. Track main until DynamicPPL 0.42 is released. Co-Authored-By: Claude Opus 4.7 (1M context) --- Project.toml | 2 +- test/Project.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 8cffa4dcf..f906480f7 100644 --- a/Project.toml +++ b/Project.toml @@ -65,4 +65,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" test = ["Pkg", "Test"] [sources] -DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "adproblems-interface"} +DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "main"} diff --git a/test/Project.toml b/test/Project.toml index 6114a8fdf..5f81fba1c 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -26,7 +26,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" [sources] -DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "adproblems-interface"} +DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "main"} [compat] ADTypes = "0.2.1, 1" From 703b8c8f49b90cc4004cf57d3011bfed80a27f1a Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 20:39:59 +0100 Subject: [PATCH 10/31] Run benchmarks and docs on Julia 1; add Julia 1 to Buildkite Pkg's [sources] field is only honoured from Julia 1.11 onward, so the LTS (1.10) jobs can't resolve the DynamicPPL main source pin. Move the benchmark and docs workflows to Julia 1, and add Julia 1 alongside 1.10 in the Buildkite CUDA matrix. Co-Authored-By: Claude Opus 4.7 (1M context) --- .buildkite/pipeline.yml | 1 + .github/workflows/Benchmark.yml | 2 +- .github/workflows/Docs.yml | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index cf5dcdf3f..0f054a6e2 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -16,3 +16,4 @@ steps: setup: julia: - "1.10" + - "1" diff --git a/.github/workflows/Benchmark.yml b/.github/workflows/Benchmark.yml index 14e7d57b6..67e53985c 100644 --- a/.github/workflows/Benchmark.yml +++ b/.github/workflows/Benchmark.yml @@ -22,7 +22,7 @@ jobs: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v2 with: - version: 'lts' + version: '1' arch: x64 - uses: actions/cache@v4 env: diff --git a/.github/workflows/Docs.yml b/.github/workflows/Docs.yml index 297438d7a..5b9d7d94e 100644 --- a/.github/workflows/Docs.yml +++ b/.github/workflows/Docs.yml @@ -25,7 +25,7 @@ jobs: - name: Build and deploy Documenter.jl docs uses: TuringLang/actions/DocsDocumenter@main with: - julia-version: 'lts' + julia-version: '1' - name: Run doctests shell: julia --project=docs --color=yes {0} From b7d820f0bdc1878ad8dab72b553e89797b241c94 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 21:48:03 +0100 Subject: [PATCH 11/31] Skip prep cache for AutoReverseDiff(; compile=true) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AbstractPPL's DI extension routes compiled-tape ReverseDiff through a closure so the tape can be reused — but that closure bakes `aref[]` in at preparation time, so mutating the context Ref between iterations silently feeds the original `aux` (stale `q_stop`, etc.) into AD and diverges the optimization. Return `nothing` from `_prepare_gradient` for `AutoReverseDiff{true}` and re-prepare in `_value_and_gradient!` so the fresh `aux` reaches AD each call without mutating shared state. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/AdvancedVI.jl | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index e2fa40ab2..e5de0bee8 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -68,6 +68,22 @@ function _value_and_gradient!( return out end +# `AutoReverseDiff(; compile=true)` is the one backend AbstractPPL's DI +# extension routes through a closure so the compiled tape can be reused. +# That closure bakes `aref[]` in at prep time, so mutating the context Ref +# between iterations would feed stale `aux` (e.g. an old `q_stop`) to AD +# and diverge the optimization. Skip caching here and re-prepare each call. +function _value_and_gradient!( + f, + out::DiffResults.MutableDiffResult, + ::Nothing, + ad::ADTypes.AutoReverseDiff{true}, + x, + aux, +) + return _value_and_gradient!(f, out, ad, x, aux) +end + """ _prepare_gradient(f, ad, x, aux) @@ -83,6 +99,10 @@ function _prepare_gradient(f, ad::ADTypes.AbstractADType, x, aux) return AbstractPPL.prepare(ad, (x, aref) -> f(x, aref[]), x; context=(Ref(aux),)) end +# See the dispatch above on `_value_and_gradient!` for the reason this returns +# `nothing` for compiled-tape ReverseDiff. +_prepare_gradient(::Any, ::ADTypes.AutoReverseDiff{true}, ::Any, ::Any) = nothing + """ restructure_ad_forward(adtype, restructure, params) From 3fdf6841e26e49777d23946d461ec2f4f1dd660c Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 21:48:10 +0100 Subject: [PATCH 12/31] Drop Zygote from benchmarks; load DI in setup blocks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AbstractPPL's DI extension provides the AD-aware `prepare` method only for backends backed by DifferentiationInterface, but doesn't ship a Zygote path. Drop Zygote from the benchmark matrix and add DifferentiationInterface so the remaining ReverseDiff/Mooncake rows resolve. Add `using DifferentiationInterface` to the @setup blocks in families.md and klminrepgraddescent.md — these blocks call AD-prepared gradients before any tutorial page loads DI. Co-Authored-By: Claude Opus 4.7 (1M context) --- bench/Project.toml | 4 ++-- bench/benchmarks.jl | 4 ++-- docs/src/families.md | 1 + docs/src/klminrepgraddescent.md | 1 + 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/bench/Project.toml b/bench/Project.toml index 4988bba16..bc0f53cad 100644 --- a/bench/Project.toml +++ b/bench/Project.toml @@ -3,6 +3,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" @@ -15,13 +16,13 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] ADTypes = "1" AdvancedVI = "0.7, 0.6" BenchmarkTools = "1" Distributions = "0.25.111" +DifferentiationInterface = "0.6, 0.7" DistributionsAD = "0.6" Enzyme = "0.13.7" FillArrays = "1" @@ -34,5 +35,4 @@ Random = "1" ReverseDiff = "1" SimpleUnPack = "1" StableRNGs = "1" -Zygote = "0.6, 0.7" julia = "1.10, 1.11.2" diff --git a/bench/benchmarks.jl b/bench/benchmarks.jl index a2819bd7d..5b3bbd0c0 100644 --- a/bench/benchmarks.jl +++ b/bench/benchmarks.jl @@ -3,7 +3,8 @@ using AdvancedVI using BenchmarkTools using Distributions using DistributionsAD -using Enzyme, ForwardDiff, ReverseDiff, Zygote, Mooncake +using DifferentiationInterface +using Enzyme, ForwardDiff, ReverseDiff, Mooncake using FillArrays using InteractiveUtils using LinearAlgebra @@ -68,7 +69,6 @@ begin ("RepGradELBO + STL", StickingTheLandingEntropy()), ], (adname, adtype) in [ - ("Zygote", AutoZygote()), ("ReverseDiff", AutoReverseDiff()), ("Mooncake", AutoMooncake(; config=Mooncake.Config())), # ("Enzyme", AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)), diff --git a/docs/src/families.md b/docs/src/families.md index a7f2e19b8..16daa27b4 100644 --- a/docs/src/families.md +++ b/docs/src/families.md @@ -139,6 +139,7 @@ using LogDensityProblems using Optimisers using Plots using ForwardDiff, ReverseDiff +using DifferentiationInterface struct Target{D} dist::D diff --git a/docs/src/klminrepgraddescent.md b/docs/src/klminrepgraddescent.md index 771cf99cd..e9c224482 100644 --- a/docs/src/klminrepgraddescent.md +++ b/docs/src/klminrepgraddescent.md @@ -104,6 +104,7 @@ using Random using Optimisers using ADTypes, ForwardDiff, ReverseDiff +using DifferentiationInterface using AdvancedVI struct Dist{D} From a57188ab82f58190a2e18b7124e9e747853ab8bc Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 22:03:51 +0100 Subject: [PATCH 13/31] Tighten comments and prefer !== for type comparisons MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Compress the compile-tape ReverseDiff rationale at the dispatch site; drop the cross-reference duplicate on `_prepare_gradient`. - Drop the stale "unlike DI's Constant" parenthetical from the DynamicPPL ext struct comment. - Add a brief WHY for the `copy(grad)` / `copy(H)` calls in `logdensity_and_gradient` / `logdensity_gradient_and_hessian` — the arrays may alias the prep's internal buffers under AbstractPPL's `!!` contract. - Use `!==` instead of `!=` when comparing type parameters in `LogDensityProblems.capabilities`. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AdvancedVIDynamicPPLExt.jl | 9 +++++---- src/AdvancedVI.jl | 9 ++------- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index 699327fc6..4363f2e3d 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -33,8 +33,7 @@ end # `model_ref`/`loglikeadj_ref` are mutated in place by `subsample`; the closure # inside `prep_grad`/`prep_hess` reads through them so the prep stays valid -# across subsampling steps (AbstractPPL bakes the closure into the prep, unlike -# DI's `Constant` which can be rebound at call time). +# across subsampling steps. # # `model_ref` is typed `Ref{Any}` because `subsample_dynamicpplmodel` returns a # `DynamicPPL.Model` whose `defaults` NamedTuple type varies with the batch — a @@ -112,6 +111,8 @@ function LogDensityProblems.logdensity(prob::DynamicPPLModelLogDensityFunction, return logdensity_impl(params, prob.model_ref[], prob.loglikeadj_ref[], prob.varinfo) end +# `!!` may alias internal buffers of `prep_*`; copy so callers can retain the +# arrays past the next AD call. function LogDensityProblems.logdensity_and_gradient( prob::DynamicPPLModelLogDensityFunction, params ) @@ -129,9 +130,9 @@ end function LogDensityProblems.capabilities( ::Type{<:DynamicPPLModelLogDensityFunction{M,L,V,ADType,PG,PH}} ) where {M,L,V,ADType<:ADTypes.AbstractADType,PG,PH} - return if PH != Nothing + return if PH !== Nothing LogDensityProblems.LogDensityOrder{2}() - elseif PG != Nothing + elseif PG !== Nothing LogDensityProblems.LogDensityOrder{1}() else LogDensityProblems.LogDensityOrder{0}() diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index e5de0bee8..cfe851afe 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -68,11 +68,8 @@ function _value_and_gradient!( return out end -# `AutoReverseDiff(; compile=true)` is the one backend AbstractPPL's DI -# extension routes through a closure so the compiled tape can be reused. -# That closure bakes `aref[]` in at prep time, so mutating the context Ref -# between iterations would feed stale `aux` (e.g. an old `q_stop`) to AD -# and diverge the optimization. Skip caching here and re-prepare each call. +# Compiled-tape ReverseDiff bakes `aref[]` into the tape at prep time, so a +# cached prep would feed stale `aux` to AD. Re-prepare each call instead. function _value_and_gradient!( f, out::DiffResults.MutableDiffResult, @@ -99,8 +96,6 @@ function _prepare_gradient(f, ad::ADTypes.AbstractADType, x, aux) return AbstractPPL.prepare(ad, (x, aref) -> f(x, aref[]), x; context=(Ref(aux),)) end -# See the dispatch above on `_value_and_gradient!` for the reason this returns -# `nothing` for compiled-tape ReverseDiff. _prepare_gradient(::Any, ::ADTypes.AutoReverseDiff{true}, ::Any, ::Any) = nothing """ From 61cee75ffbca1d88c7eed718f424051ba01f3b2b Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 22:03:56 +0100 Subject: [PATCH 14/31] Drop ReverseDiff from benchmarks Mooncake covers the reverse-mode benchmark slot and is the backend we actually want to track; ReverseDiff adds noise without informing the gradient-time numbers we report. Co-Authored-By: Claude Opus 4.7 (1M context) --- bench/Project.toml | 2 -- bench/benchmarks.jl | 3 +-- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/bench/Project.toml b/bench/Project.toml index bc0f53cad..68e775d0c 100644 --- a/bench/Project.toml +++ b/bench/Project.toml @@ -13,7 +13,6 @@ LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" @@ -32,7 +31,6 @@ LogDensityProblems = "2" Mooncake = "0.4.5, 0.5" Optimisers = "0.3, 0.4" Random = "1" -ReverseDiff = "1" SimpleUnPack = "1" StableRNGs = "1" julia = "1.10, 1.11.2" diff --git a/bench/benchmarks.jl b/bench/benchmarks.jl index 5b3bbd0c0..e7ce5a346 100644 --- a/bench/benchmarks.jl +++ b/bench/benchmarks.jl @@ -4,7 +4,7 @@ using BenchmarkTools using Distributions using DistributionsAD using DifferentiationInterface -using Enzyme, ForwardDiff, ReverseDiff, Mooncake +using Enzyme, ForwardDiff, Mooncake using FillArrays using InteractiveUtils using LinearAlgebra @@ -69,7 +69,6 @@ begin ("RepGradELBO + STL", StickingTheLandingEntropy()), ], (adname, adtype) in [ - ("ReverseDiff", AutoReverseDiff()), ("Mooncake", AutoMooncake(; config=Mooncake.Config())), # ("Enzyme", AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)), ], From 4a170ea36c7f11750f4eb912293b0f5f85f94ba2 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 22:14:45 +0100 Subject: [PATCH 15/31] Apply JuliaFormatter suggestion in bench/benchmarks.jl Collapse the single-row AD list onto the opening bracket per the formatter check on PR #255. Co-Authored-By: Claude Opus 4.7 (1M context) --- bench/benchmarks.jl | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/bench/benchmarks.jl b/bench/benchmarks.jl index e7ce5a346..888137d20 100644 --- a/bench/benchmarks.jl +++ b/bench/benchmarks.jl @@ -68,9 +68,8 @@ begin ("RepGradELBO", ClosedFormEntropy()), ("RepGradELBO + STL", StickingTheLandingEntropy()), ], - (adname, adtype) in [ - ("Mooncake", AutoMooncake(; config=Mooncake.Config())), - # ("Enzyme", AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)), + (adname, adtype) in [("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)))), From 53a672fbcf75c5ee9d522afa7329b9dcdb220c5e Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 26 May 2026 10:27:41 +0100 Subject: [PATCH 16/31] Drop DynamicPPL main source pin DynamicPPL has a new release covering the AbstractPPL/DynamicPPL migration, so the source override is no longer needed. Co-Authored-By: Claude Opus 4.7 (1M context) --- Project.toml | 3 --- test/Project.toml | 3 --- 2 files changed, 6 deletions(-) diff --git a/Project.toml b/Project.toml index f906480f7..7fe21d671 100644 --- a/Project.toml +++ b/Project.toml @@ -63,6 +63,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] test = ["Pkg", "Test"] - -[sources] -DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "main"} diff --git a/test/Project.toml b/test/Project.toml index 5f81fba1c..55acfbe31 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -25,9 +25,6 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" -[sources] -DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "main"} - [compat] ADTypes = "0.2.1, 1" Bijectors = "0.16" From 7f24addfc565a6882090bc1d82b52cc27fabc660 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 26 May 2026 11:29:53 +0100 Subject: [PATCH 17/31] Switch docs and benchmarks to Mooncake MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Drop DifferentiationInterface from bench and docs (DI is unused in AdvancedVI itself, which routes AD through AbstractPPL). Keep it in test/ where it remains load-bearing as the trigger for DI's Enzyme extension. - Replace ReverseDiff with Mooncake everywhere in docs and benchmarks; use the no-arg `AutoMooncake()` form consistently. - In `MvLocationScaleLowRank.entropy`, swap `Hermitian` for `Symmetric` so that Mooncake's `logdet ∘ Symmetric` rule (chalk-lab/Mooncake.jl PR #1055) fires — `Hermitian` is not covered by that rule and falls through to a missing `dsytrf` foreigncall. The two wrappers are semantically identical for real element types. - In the QMC tutorial section, declare `QuasiMonteCarlo.sample` as non-differentiable via `Mooncake.@zero_derivative` so Mooncake skips the Sobol/Owen bit-twiddling that would otherwise trigger its bitcast safety check. - Add `ENV["GKSwstype"] = "100"` at the top of `docs/make.jl` so GR uses the no-display workstation; `savefig` still writes files but no interactive plot window is opened during the build. Co-Authored-By: Claude Opus 4.7 (1M context) --- bench/Project.toml | 2 -- bench/benchmarks.jl | 6 +++--- docs/Project.toml | 6 ++---- docs/make.jl | 5 +++++ docs/src/families.md | 5 ++--- docs/src/klminrepgraddescent.md | 13 ++++++++----- docs/src/tutorials/basic.md | 5 ++--- docs/src/tutorials/constrained.md | 4 ++-- docs/src/tutorials/flows.md | 7 +++---- docs/src/tutorials/stan.md | 4 ++-- docs/src/tutorials/subsampling.md | 8 ++++---- src/families/location_scale_low_rank.jl | 3 ++- test/runtests.jl | 4 ++-- 13 files changed, 37 insertions(+), 35 deletions(-) diff --git a/bench/Project.toml b/bench/Project.toml index 68e775d0c..1d3e954a1 100644 --- a/bench/Project.toml +++ b/bench/Project.toml @@ -3,7 +3,6 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" @@ -21,7 +20,6 @@ ADTypes = "1" AdvancedVI = "0.7, 0.6" BenchmarkTools = "1" Distributions = "0.25.111" -DifferentiationInterface = "0.6, 0.7" DistributionsAD = "0.6" Enzyme = "0.13.7" FillArrays = "1" diff --git a/bench/benchmarks.jl b/bench/benchmarks.jl index 888137d20..878268faf 100644 --- a/bench/benchmarks.jl +++ b/bench/benchmarks.jl @@ -3,7 +3,6 @@ using AdvancedVI using BenchmarkTools using Distributions using DistributionsAD -using DifferentiationInterface using Enzyme, ForwardDiff, Mooncake using FillArrays using InteractiveUtils @@ -68,8 +67,9 @@ begin ("RepGradELBO", ClosedFormEntropy()), ("RepGradELBO + STL", StickingTheLandingEntropy()), ], - (adname, adtype) in [("Mooncake", AutoMooncake(; config=Mooncake.Config())), - # ("Enzyme", AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)), + (adname, adtype) in [ + ("Mooncake", AutoMooncake()), + # ("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)))), diff --git a/docs/Project.toml b/docs/Project.toml index 9374d4c52..d8793831a 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -4,7 +4,6 @@ Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" @@ -13,12 +12,12 @@ Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256" OpenML = "8b6db2d4-7670-4922-a472-f9537c81ab66" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b" -ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" StanLogDensityProblems = "a545de4d-8dba-46db-9d34-4e41d3f07807" StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c" @@ -28,7 +27,6 @@ Accessors = "0.1" AdvancedVI = "0.7, 0.6" Bijectors = "0.15, 0.16" DataFrames = "1" -DifferentiationInterface = "0.7" Distributions = "0.25" Documenter = "1" FillArrays = "1" @@ -37,12 +35,12 @@ Functors = "0.5" JSON = "0.21, 1" LogDensityProblems = "2.1.1" LogDensityProblemsAD = "1" +Mooncake = "0.4, 0.5" NormalizingFlows = "0.2.2" OpenML = "0.3" Optimisers = "0.3, 0.4" Plots = "1" QuasiMonteCarlo = "0.3" -ReverseDiff = "1" StanLogDensityProblems = "0.1" StatsFuns = "1" julia = "1.10, 1.11.2" diff --git a/docs/make.jl b/docs/make.jl index f85540832..44acd186d 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,4 +1,9 @@ +# Force GR (Plots.jl's default backend) into its no-display workstation so +# `savefig` writes files but no interactive window is opened. Must be set +# before GR loads. +ENV["GKSwstype"] = "100" + using AdvancedVI using Documenter diff --git a/docs/src/families.md b/docs/src/families.md index 16daa27b4..57dee3333 100644 --- a/docs/src/families.md +++ b/docs/src/families.md @@ -138,8 +138,7 @@ using LinearAlgebra using LogDensityProblems using Optimisers using Plots -using ForwardDiff, ReverseDiff -using DifferentiationInterface +using ForwardDiff, Mooncake struct Target{D} dist::D @@ -185,7 +184,7 @@ D = ones(n_dims) U = zeros(n_dims, 3) q0_lr = LowRankGaussian(μ, D, U) -alg = KLMinRepGradDescent(AutoReverseDiff(); optimizer=Adam(0.01), operator=ClipScale()) +alg = KLMinRepGradDescent(AutoMooncake(); optimizer=Adam(0.01), operator=ClipScale()) max_iter = 10^4 diff --git a/docs/src/klminrepgraddescent.md b/docs/src/klminrepgraddescent.md index e9c224482..7fd22abbc 100644 --- a/docs/src/klminrepgraddescent.md +++ b/docs/src/klminrepgraddescent.md @@ -103,8 +103,7 @@ using Plots using Random using Optimisers -using ADTypes, ForwardDiff, ReverseDiff -using DifferentiationInterface +using ADTypes, ForwardDiff, Mooncake using AdvancedVI struct Dist{D} @@ -151,7 +150,7 @@ Recall that the original ADVI objective with a closed-form entropy (CFE) is give n_montecarlo = 16; cfe = KLMinRepGradDescent( - AutoReverseDiff(); + AutoMooncake(); entropy=ClosedFormEntropy(), optimizer=Adam(1e-2), operator=ClipScale(), @@ -163,7 +162,7 @@ The repgradelbo estimator can instead be created as follows: ```@example repgradelbo stl = KLMinRepGradDescent( - AutoReverseDiff(); + AutoMooncake(); entropy=StickingTheLandingEntropy(), optimizer=Adam(1e-2), operator=ClipScale(), @@ -263,6 +262,10 @@ In this case, it suffices to override its `rand` specialization as follows: using QuasiMonteCarlo using StatsFuns +# QMC samples are inputs to AD, not parameters; declare them non-differentiable +# so Mooncake doesn't trace through Sobol/Owen bit-twiddling. +Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(QuasiMonteCarlo.sample), Vararg} + qmcrng = SobolSample(; R=OwenScramble(; base=2, pad=32)) function Distributions.rand( @@ -282,7 +285,7 @@ nothing ```@setup repgradelbo _, info_qmc, _ = AdvancedVI.optimize( - KLMinRepGradDescent(AutoReverseDiff(); n_samples=n_montecarlo, optimizer=Adam(1e-2), operator=ClipScale()), + KLMinRepGradDescent(AutoMooncake(); n_samples=n_montecarlo, optimizer=Adam(1e-2), operator=ClipScale()), max_iter, model, q0; diff --git a/docs/src/tutorials/basic.md b/docs/src/tutorials/basic.md index 8ed6de6bd..450de5414 100644 --- a/docs/src/tutorials/basic.md +++ b/docs/src/tutorials/basic.md @@ -143,10 +143,10 @@ nothing For the VI algorithm, we will use `KLMinRepGradDescent`: ```@example basic -using ADTypes, ReverseDiff +using ADTypes, Mooncake using AdvancedVI -alg = KLMinRepGradDescent(ADTypes.AutoReverseDiff(); operator=ClipScale()); +alg = KLMinRepGradDescent(ADTypes.AutoMooncake(); operator=ClipScale()); nothing ``` @@ -166,7 +166,6 @@ For this example, we will use `LogDensityProblemsAD` to equip our problem with a [^RMW2014]: Rezende, D. J., Mohamed, S., & Wierstra, D. (2014, June). Stochastic backpropagation and approximate inference in deep generative models. In *International Conference on Machine Learning*. PMLR. [^KW2014]: Kingma, D. P., & Welling, M. (2014). Auto-encoding variational bayes. In *International Conference on Learning Representations*. ```@example basic -using DifferentiationInterface: DifferentiationInterface using LogDensityProblemsAD: LogDensityProblemsAD prob_trans_ad = LogDensityProblemsAD.ADgradient(ADTypes.AutoForwardDiff(), prob_trans) diff --git a/docs/src/tutorials/constrained.md b/docs/src/tutorials/constrained.md index 857e6a3af..6f78019fb 100644 --- a/docs/src/tutorials/constrained.md +++ b/docs/src/tutorials/constrained.md @@ -192,10 +192,10 @@ We can also wrap `prob_trans` with `LogDensityProblemsAD.ADGradient` to make it ```@example constraints using LogDensityProblemsAD -using ADTypes, ReverseDiff +using ADTypes, Mooncake prob_trans_ad = LogDensityProblemsAD.ADgradient( - ADTypes.AutoReverseDiff(; compile=true), prob_trans; x=randn(2) + ADTypes.AutoMooncake(), prob_trans; x=randn(2) ) ``` diff --git a/docs/src/tutorials/flows.md b/docs/src/tutorials/flows.md index a226f36f6..2fcf5bd4d 100644 --- a/docs/src/tutorials/flows.md +++ b/docs/src/tutorials/flows.md @@ -138,8 +138,7 @@ Let's instantiate the model: ```@example flow using ADTypes: ADTypes -using ReverseDiff: ReverseDiff -using DifferentiationInterface: DifferentiationInterface +using Mooncake: Mooncake using LogDensityProblemsAD: LogDensityProblemsAD prob_trans = TransformedLogDensityProblem(prob) @@ -159,7 +158,7 @@ d = LogDensityProblems.dimension(prob_trans_ad) q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(I, d, d))) max_iter = 3*10^3 -alg = KLMinRepGradProxDescent(ADTypes.AutoReverseDiff(; compile=true)) +alg = KLMinRepGradProxDescent(ADTypes.AutoMooncake()) q_out, info, _ = AdvancedVI.optimize(alg, max_iter, prob_trans_ad, q; show_progress=false) b = Bijectors.bijector(prob) binv = Bijectors.inverse(b) @@ -229,7 +228,7 @@ Furthermore, Agrawal *et al.*[^AD2025] claim that using a larger number of Monte using Optimisers: Optimisers alg_flow = KLMinRepGradDescent( - ADTypes.AutoReverseDiff(; compile=true); + ADTypes.AutoMooncake(); n_samples=8, optimizer=Optimisers.Adam(1e-2), operator=IdentityOperator(), diff --git a/docs/src/tutorials/stan.md b/docs/src/tutorials/stan.md index 1890675d6..2bfa80997 100644 --- a/docs/src/tutorials/stan.md +++ b/docs/src/tutorials/stan.md @@ -75,13 +75,13 @@ nothing The rest is the same as all `LogDensityProblem` with the exception of how to deal with constrainted variables: Since `StanLogDensityProblems` automatically transforms the support of the target problem to be unconstrained, we do not need to involve `Bijectors`. ```@example stan -using ADTypes, ReverseDiff +using ADTypes, Mooncake using AdvancedVI using LinearAlgebra using LogDensityProblems using Plots -alg = KLMinRepGradDescent(ADTypes.AutoReverseDiff(); operator=ClipScale()) +alg = KLMinRepGradDescent(ADTypes.AutoMooncake(); operator=ClipScale()) d = LogDensityProblems.dimension(model) q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.37*I, d, d))) diff --git a/docs/src/tutorials/subsampling.md b/docs/src/tutorials/subsampling.md index b2b28b316..0bdc650dc 100644 --- a/docs/src/tutorials/subsampling.md +++ b/docs/src/tutorials/subsampling.md @@ -83,12 +83,12 @@ nothing Let's now istantiate the model and set up automatic differentiation using [`LogDensityProblemsAD`](https://github.com/tpapp/LogDensityProblemsAD.jl?tab=readme-ov-file). ```@example subsampling -using ADTypes, ReverseDiff +using ADTypes, Mooncake using LogDensityProblemsAD prob = LogReg(X, y, size(X, 1)) prob_ad = LogDensityProblemsAD.ADgradient( - ADTypes.AutoReverseDiff(), prob; x=randn(LogDensityProblems.dimension(prob)) + ADTypes.AutoMooncake(), prob; x=randn(LogDensityProblems.dimension(prob)) ) nothing ``` @@ -129,7 +129,7 @@ We will us a batch size of 32, which results in `313 = length(subsampling) = cei dataset = 1:size(prob.X, 1) batchsize = 32 subsampling = ReshufflingBatchSubsampling(dataset, batchsize) -alg_sub = KLMinRepGradProxDescent(ADTypes.AutoReverseDiff(; compile=true); subsampling) +alg_sub = KLMinRepGradProxDescent(ADTypes.AutoMooncake(); subsampling) nothing ``` @@ -148,7 +148,7 @@ nothing If we don't supply a subsampling strategy to `KLMinRepGradProxDescent`, subsampling will not be used. ```@example subsampling -alg_full = KLMinRepGradProxDescent(ADTypes.AutoReverseDiff(; compile=true)) +alg_full = KLMinRepGradProxDescent(ADTypes.AutoMooncake()) nothing ``` diff --git a/src/families/location_scale_low_rank.jl b/src/families/location_scale_low_rank.jl index f3e444a30..6703bcea3 100644 --- a/src/families/location_scale_low_rank.jl +++ b/src/families/location_scale_low_rank.jl @@ -36,7 +36,8 @@ function StatsBase.entropy(q::MvLocationScaleLowRank) (; location, scale_diag, scale_factors, dist) = q n_dims = length(location) scale_diag2 = scale_diag .* scale_diag - UtDinvU = Hermitian(scale_factors' * (scale_factors ./ scale_diag2)) + # `Symmetric` rather than `Hermitian`: Mooncake lacks a rule for the `Hermitian` path. + UtDinvU = Symmetric(scale_factors' * (scale_factors ./ scale_diag2)) logdetΣ = 2 * sum(log.(scale_diag)) + logdet(I + UtDinvU) return n_dims * convert(eltype(location), entropy(dist)) + logdetΣ / 2 end diff --git a/test/runtests.jl b/test/runtests.jl index 519542e99..464d7e1c9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -17,7 +17,7 @@ using Random, StableRNGs using Statistics using StatsBase -using DifferentiationInterface +using DifferentiationInterface # triggers DI's Enzyme extension using AdvancedVI const PROGRESS = haskey(ENV, "PROGRESS") @@ -30,7 +30,7 @@ elseif AD_str == "ForwardDiff" AutoForwardDiff() elseif AD_str == "Mooncake" using Mooncake - AutoMooncake(; config=Mooncake.Config()) + AutoMooncake() elseif AD_str == "Enzyme" using Enzyme AutoEnzyme(; From 4e9693ead0bb4ff3ddffb8a3fd09c2417d8ba2bb Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 26 May 2026 12:28:38 +0100 Subject: [PATCH 18/31] Switch format workflow to TuringLang/actions/Format Use the shared TuringLang Format action (JuliaFormatter v1, default paths, suggestions disabled) instead of the custom reviewdog pipeline. Matches the AbstractPPL.jl workflow. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/Format.yml | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/.github/workflows/Format.yml b/.github/workflows/Format.yml index 732212c62..920bf1c02 100644 --- a/.github/workflows/Format.yml +++ b/.github/workflows/Format.yml @@ -1,8 +1,11 @@ -name: Format suggestions +name: Format on: + push: + branches: + - main pull_request: - + concurrency: # Skip intermediate builds: always. # Cancel intermediate builds: only if it is a pull request build. @@ -12,15 +15,7 @@ concurrency: jobs: format: runs-on: ubuntu-latest + steps: - - uses: actions/checkout@v4 - - uses: julia-actions/setup-julia@v2 - with: - version: 1 - - run: | - julia -e 'using Pkg; Pkg.add("JuliaFormatter")' - julia -e 'using JuliaFormatter; format("."; verbose=true)' - - uses: reviewdog/action-suggester@v1 - with: - tool_name: JuliaFormatter - fail_on_error: true + - name: Format code + uses: TuringLang/actions/Format@main From 773460c99df00d779e8a23d5fd2b97b542f209fe Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 26 May 2026 12:28:54 +0100 Subject: [PATCH 19/31] Drop LogDensityProblemsAD from docs tutorials `LogDensityProblemsAD.ADgradient` has no method for `AutoMooncake`, so the previous wrapping pattern broke once the docs migrated to Mooncake. `AdvancedVI.optimize` already differentiates through `LogDensityProblems.logdensity` using the algorithm's `adtype`, so the wrapper is no longer needed: - basic, flows, subsampling: pass `prob_trans`/`prob` directly to `AdvancedVI.optimize`. In subsampling, specialize `AdvancedVI.subsample` on `LogReg` rather than on the (now-removed) `ADgradient` wrapper. - constrained: `FisherMinBatchMatch` still requires order-1 capability, so give `TransformedLogDensityProblem` its own `logdensity_and_gradient` via `AbstractPPL.prepare(AutoForwardDiff(), ...)` and bump `capabilities` to `LogDensityOrder{1}`. - docs/Project.toml: drop `LogDensityProblemsAD`, add `AbstractPPL`. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/Project.toml | 4 ++-- docs/src/tutorials/basic.md | 27 ++++++++-------------- docs/src/tutorials/constrained.md | 37 +++++++++++++++---------------- docs/src/tutorials/flows.md | 12 ++++------ docs/src/tutorials/subsampling.md | 35 +++++++++++------------------ 5 files changed, 46 insertions(+), 69 deletions(-) diff --git a/docs/Project.toml b/docs/Project.toml index d8793831a..ee4daaa3f 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,5 +1,6 @@ [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" @@ -11,7 +12,6 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" -LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256" OpenML = "8b6db2d4-7670-4922-a472-f9537c81ab66" @@ -23,6 +23,7 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c" [compat] ADTypes = "1" +AbstractPPL = "0.15" Accessors = "0.1" AdvancedVI = "0.7, 0.6" Bijectors = "0.15, 0.16" @@ -34,7 +35,6 @@ ForwardDiff = "0.10, 1" Functors = "0.5" JSON = "0.21, 1" LogDensityProblems = "2.1.1" -LogDensityProblemsAD = "1" Mooncake = "0.4, 0.5" NormalizingFlows = "0.2.2" OpenML = "0.3" diff --git a/docs/src/tutorials/basic.md b/docs/src/tutorials/basic.md index 450de5414..ee1b8ad81 100644 --- a/docs/src/tutorials/basic.md +++ b/docs/src/tutorials/basic.md @@ -35,7 +35,7 @@ function LogDensityProblems.logdensity(model::LogReg, θ) logprior_β = logpdf(MvNormal(Zeros(d), σ), β) logprior_σ = logpdf(LogNormal(0, 3), σ) - logit = X*β + logit = X * β loglike_y = mapreduce((li, yi) -> logpdf(BernoulliLogit(li), yi), +, logit, y) return loglike_y + logprior_β + logprior_σ end @@ -159,26 +159,17 @@ Location-scale family distributions require the scale matrix to have strictly po Here, the projection operator `ClipScale` ensures this. `KLMinRepGradDescent`, in particular, assumes that the target `LogDensityProblem` is differentiable. -If the `LogDensityProblem` has a differentiation [capability](https://www.tamaspapp.eu/LogDensityProblems.jl/dev/#LogDensityProblems.capabilities) of at least first-order, we can take advantage of this. -For this example, we will use `LogDensityProblemsAD` to equip our problem with a first-order capability: +If the problem provides only a zeroth-order [capability](https://www.tamaspapp.eu/LogDensityProblems.jl/dev/#LogDensityProblems.capabilities), `AdvancedVI` will differentiate through `LogDensityProblems.logdensity` directly using the AD backend supplied to the algorithm. [^TL2014]: Titsias, M., & Lázaro-Gredilla, M. (2014, June). Doubly stochastic variational Bayes for non-conjugate inference. In *International Conference on Machine Learning*. PMLR. [^RMW2014]: Rezende, D. J., Mohamed, S., & Wierstra, D. (2014, June). Stochastic backpropagation and approximate inference in deep generative models. In *International Conference on Machine Learning*. PMLR. [^KW2014]: Kingma, D. P., & Welling, M. (2014). Auto-encoding variational bayes. In *International Conference on Learning Representations*. -```@example basic -using LogDensityProblemsAD: LogDensityProblemsAD - -prob_trans_ad = LogDensityProblemsAD.ADgradient(ADTypes.AutoForwardDiff(), prob_trans) -nothing -``` - -For the variational family, we will consider a `FullRankGaussian` approximation: - + For the variational family, we will consider a `FullRankGaussian` approximation: ```@example basic using LinearAlgebra d = LogDensityProblems.dimension(prob_trans) -q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6*I, d, d))) +q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6 * I, d, d))) nothing ``` @@ -189,7 +180,7 @@ We can now run VI: ```@example basic max_iter = 10^4 -q_out, info, _ = AdvancedVI.optimize(alg, max_iter, prob_trans_ad, q; show_progress=false) +q_out, info, _ = AdvancedVI.optimize(alg, max_iter, prob_trans, q; show_progress=false) nothing ``` @@ -244,9 +235,9 @@ using StatsFuns: StatsFuns Approximate the posterior predictive probability for a logistic link function using Mackay's approximation (Bishop p. 220). """ function logistic_prediction(X, μ_β, Σ_β) - xtΣx = sum((prob.X*Σ_β) .* prob.X; dims=2)[:, 1] - κ = @. 1/sqrt(1 + π/8*xtΣx) - return StatsFuns.logistic.(κ .* X*μ_β) + xtΣx = sum((prob.X * Σ_β) .* prob.X; dims=2)[:, 1] + κ = @. 1 / sqrt(1 + π / 8 * xtΣx) + return StatsFuns.logistic.(κ .* X * μ_β) end logging_interval = 100 @@ -285,7 +276,7 @@ The `callback` can be supplied to `optimize`: ```@example basic max_iter = 10^4 q_out, info, _ = AdvancedVI.optimize( - alg, max_iter, prob_trans_ad, q; show_progress=false, callback=callback + alg, max_iter, prob_trans, q; show_progress=false, callback=callback ) nothing ``` diff --git a/docs/src/tutorials/constrained.md b/docs/src/tutorials/constrained.md index 6f78019fb..73a060902 100644 --- a/docs/src/tutorials/constrained.md +++ b/docs/src/tutorials/constrained.md @@ -147,6 +147,10 @@ This approach only requires the user to implement the model-specific `Bijectors. The rest can be done by simply copy-pasting the code below: ```@example constraints +using ADTypes +using AbstractPPL: AbstractPPL +using ForwardDiff: ForwardDiff + struct TransformedLogDensityProblem{Prob,BInv} prob::Prob binv::BInv @@ -164,6 +168,14 @@ function LogDensityProblems.logdensity(prob_trans::TransformedLogDensityProblem, return LogDensityProblems.logdensity(prob, θ) + logabsdetjac end +function LogDensityProblems.logdensity_and_gradient( + prob_trans::TransformedLogDensityProblem, θ_trans +) + f = Base.Fix1(LogDensityProblems.logdensity, prob_trans) + prep = AbstractPPL.prepare(AutoForwardDiff(), f, θ_trans) + return AbstractPPL.value_and_gradient!!(prep, θ_trans) +end + function LogDensityProblems.dimension(prob_trans::TransformedLogDensityProblem) (; prob, binv) = prob_trans b = Bijectors.inverse(binv) @@ -171,10 +183,8 @@ function LogDensityProblems.dimension(prob_trans::TransformedLogDensityProblem) return prod(Bijectors.output_size(b, (d,))) end -function LogDensityProblems.capabilities( - ::Type{TransformedLogDensityProblem{Prob,BInv}} -) where {Prob,BInv} - return LogDensityProblems.capabilities(Prob) +function LogDensityProblems.capabilities(::Type{<:TransformedLogDensityProblem}) + return LogDensityProblems.LogDensityOrder{1}() end nothing ``` @@ -188,17 +198,6 @@ x = randn(LogDensityProblems.dimension(prob_trans)) # sample on an unconstrained LogDensityProblems.logdensity(prob_trans, x) ``` -We can also wrap `prob_trans` with `LogDensityProblemsAD.ADGradient` to make it differentiable. - -```@example constraints -using LogDensityProblemsAD -using ADTypes, Mooncake - -prob_trans_ad = LogDensityProblemsAD.ADgradient( - ADTypes.AutoMooncake(), prob_trans; x=randn(2) -) -``` - Let's now run VI to verify that it works. Here, we will use `FisherMinBatchMatch`, which expects an unconstrained posterior. @@ -206,11 +205,11 @@ Here, we will use `FisherMinBatchMatch`, which expects an unconstrained posterio using AdvancedVI using LinearAlgebra -d = LogDensityProblems.dimension(prob_trans_ad) -q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6*I, d, d))) +d = LogDensityProblems.dimension(prob_trans) +q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6 * I, d, d))) q_opt, info, _ = AdvancedVI.optimize( - FisherMinBatchMatch(), 100, prob_trans_ad, q; show_progress=false + FisherMinBatchMatch(), 100, prob_trans, q; show_progress=false ) nothing ``` @@ -270,7 +269,7 @@ end LogDensityProblems.dimension(::MeanTransformed) = 2 function LogDensityProblems.capabilities(::Type{MeanTransformed}) - LogDensityProblems.LogDensityOrder{0}() + return LogDensityProblems.LogDensityOrder{0}() end n_data = 30 diff --git a/docs/src/tutorials/flows.md b/docs/src/tutorials/flows.md index 2fcf5bd4d..4da07546c 100644 --- a/docs/src/tutorials/flows.md +++ b/docs/src/tutorials/flows.md @@ -139,12 +139,8 @@ Let's instantiate the model: ```@example flow using ADTypes: ADTypes using Mooncake: Mooncake -using LogDensityProblemsAD: LogDensityProblemsAD prob_trans = TransformedLogDensityProblem(prob) -prob_trans_ad = LogDensityProblemsAD.ADgradient( - ADTypes.AutoForwardDiff(), prob_trans; x=[1.0, 1.0] -) nothing ``` @@ -154,12 +150,12 @@ For the algorithm, we will use the `KLMinRepGradProxDescent` objective. using AdvancedVI using LinearAlgebra -d = LogDensityProblems.dimension(prob_trans_ad) +d = LogDensityProblems.dimension(prob_trans) q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(I, d, d))) -max_iter = 3*10^3 +max_iter = 3 * 10^3 alg = KLMinRepGradProxDescent(ADTypes.AutoMooncake()) -q_out, info, _ = AdvancedVI.optimize(alg, max_iter, prob_trans_ad, q; show_progress=false) +q_out, info, _ = AdvancedVI.optimize(alg, max_iter, prob_trans, q; show_progress=false) b = Bijectors.bijector(prob) binv = Bijectors.inverse(b) q_out_trans = Bijectors.TransformedDistribution(q_out, binv) @@ -242,7 +238,7 @@ Without further due, let's now run VI: ```@example flow max_iter = 300 q_flow_out, info_flow, _ = AdvancedVI.optimize( - alg_flow, max_iter, prob_trans_ad, q_flow; show_progress=false + alg_flow, max_iter, prob_trans, q_flow; show_progress=false ) q_flow_out_trans = Bijectors.TransformedDistribution(q_flow_out, binv) nothing diff --git a/docs/src/tutorials/subsampling.md b/docs/src/tutorials/subsampling.md index 0bdc650dc..294781e63 100644 --- a/docs/src/tutorials/subsampling.md +++ b/docs/src/tutorials/subsampling.md @@ -32,9 +32,9 @@ function LogDensityProblems.logdensity(model::LogReg, θ) logprior_β = logpdf(MvNormal(Zeros(d), σ), β) logprior_σ = logpdf(Normal(0, 3), σ) - logit = X*β + logit = X * β loglike_y = mapreduce((li, yi) -> logpdf(BernoulliLogit(li), yi), +, logit, y) - return n_data/n*loglike_y + logprior_β + logprior_σ + return n_data / n * loglike_y + logprior_β + logprior_σ end function LogDensityProblems.dimension(model::LogReg) @@ -80,34 +80,25 @@ X = hcat(X, ones(size(X, 1))) nothing ``` -Let's now istantiate the model and set up automatic differentiation using [`LogDensityProblemsAD`](https://github.com/tpapp/LogDensityProblemsAD.jl?tab=readme-ov-file). +Let's now instantiate the model. ```@example subsampling using ADTypes, Mooncake -using LogDensityProblemsAD prob = LogReg(X, y, size(X, 1)) -prob_ad = LogDensityProblemsAD.ADgradient( - ADTypes.AutoMooncake(), prob; x=randn(LogDensityProblems.dimension(prob)) -) nothing ``` To enable subsampling, `LogReg` has to implement the method `AdvancedVI.subsample`. For our model, this is fairly simple: We only need to select the rows of `X` and the elements of `y` corresponding to the batch of data points. -As subtle point here is that we wrapped `prob` with `LogDensityProblemsAD.ADgradient` into `prob_ad`. -Therefore, `AdvancedVI` sees `prob_ad` and not `prob`. -This means we have to specialize `AdvancedVI.subsample` to `typeof(prob_ad)` and not `LogReg`. ```@example subsampling using Accessors using AdvancedVI -function AdvancedVI.subsample(prob::typeof(prob_ad), idx) - (; X, y, n_data) = parent(prob_ad) - prob′ = @set prob.ℓ.X = X[idx, :] - prob′′ = @set prob′.ℓ.y = y[idx] - return prob′′ +function AdvancedVI.subsample(prob::LogReg, idx) + prob′ = @set prob.X = prob.X[idx, :] + return @set prob′.y = prob′.y[idx] end nothing ``` @@ -157,8 +148,8 @@ The variational family will be set up as follows: ```@example subsampling using LinearAlgebra -d = LogDensityProblems.dimension(prob_ad) -q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6*I, d, d))) +d = LogDensityProblems.dimension(prob) +q = FullRankGaussian(zeros(d), LowerTriangular(Matrix{Float64}(0.6 * I, d, d))) nothing ``` @@ -177,9 +168,9 @@ time_begin = nothing Approximate the posterior predictive probability for a logistic link function using Mackay's approximation (Bishop p. 220). """ function logistic_prediction(X, μ_β, Σ_β) - xtΣx = sum((prob.X*Σ_β) .* prob.X; dims=2)[:, 1] - κ = @. 1/sqrt(1 + π/8*xtΣx) - return StatsFuns.logistic.(κ .* X*μ_β) + xtΣx = sum((prob.X * Σ_β) .* prob.X; dims=2)[:, 1] + κ = @. 1 / sqrt(1 + π / 8 * xtΣx) + return StatsFuns.logistic.(κ .* X * μ_β) end function callback(; iteration, averaged_params, restructure, kwargs...) @@ -208,12 +199,12 @@ end time_begin = time() _, info_full, _ = AdvancedVI.optimize( - alg_full, max_iter, prob_ad, q; show_progress=false, callback + alg_full, max_iter, prob, q; show_progress=false, callback ); time_begin = time() _, info_sub, _ = AdvancedVI.optimize( - alg_sub, max_iter, prob_ad, q; show_progress=false, callback + alg_sub, max_iter, prob, q; show_progress=false, callback ); nothing ``` From 9c2e44e3d04f562987abeb34f73bc9691a1a5f61 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 26 May 2026 12:29:03 +0100 Subject: [PATCH 20/31] Move Mooncake overlay for Bijectors.Stacked into basic.md MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Loading `OpenML` transitively brings in `CategoricalArrays`, which adds `Base.similar` methods with unbound type parameters. Those break Mooncake's compilation of `collect(::Base.Generator)` — used by `Bijectors.Stacked`'s `_with_logabsdet_jacobian` via `map(zip(...)) do ...`. Place the `Mooncake.@mooncake_overlay` workaround in basic.md (the first tutorial that triggers the path) inside a `@setup` block so it runs invisibly during the build, and surface the explanation + listing in a collapsible `
` block so readers can see what it does without the noise being inline. The overlay extends `Mooncake.mooncake_method_table` globally, so once basic.md runs the override is in effect for all subsequent tutorials. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/src/tutorials/basic.md | 51 +++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/docs/src/tutorials/basic.md b/docs/src/tutorials/basic.md index ee1b8ad81..ecf278289 100644 --- a/docs/src/tutorials/basic.md +++ b/docs/src/tutorials/basic.md @@ -120,6 +120,57 @@ y = Vector{Bool}(data[:, end] .== "Mine") nothing ``` +```@setup basic +using Bijectors: Bijectors +using Mooncake: Mooncake + +Mooncake.@mooncake_overlay function Bijectors._with_logabsdet_jacobian( + sb::Bijectors.Stacked, x::AbstractVector +) + T = eltype(x) + y = T[] + logjac = zero(T) + for i in eachindex(sb.bs) + yi, lji = Bijectors.with_logabsdet_jacobian(sb.bs[i], x[sb.ranges_in[i]]) + y = vcat(y, yi) + logjac = logjac + sum(lji) + end + return (y, logjac) +end +``` + +```@raw html +
+Why does Mooncake need help differentiating Bijectors.Stacked when OpenML is loaded? (click to expand) +``` + +Loading `OpenML` transitively brings in `CategoricalArrays`, which adds +`Base.similar` methods with unbound type parameters. Those break Mooncake's +compilation of `collect(::Base.Generator)` — used by `Bijectors.Stacked`'s +`_with_logabsdet_jacobian` via `map(zip(...)) do ...`. The setup block above +sidesteps it with `Mooncake.@mooncake_overlay`, replacing the implementation +with an explicit loop: + +```julia +Mooncake.@mooncake_overlay function Bijectors._with_logabsdet_jacobian( + sb::Bijectors.Stacked, x::AbstractVector +) + T = eltype(x) + y = T[] + logjac = zero(T) + for i in eachindex(sb.bs) + yi, lji = Bijectors.with_logabsdet_jacobian(sb.bs[i], x[sb.ranges_in[i]]) + y = vcat(y, yi) + logjac = logjac + sum(lji) + end + return (y, logjac) +end +``` + +```@raw html +
+``` + Let's apply some basic pre-processing and add an intercept column: ```@example basic From 2fbbcdc24011e8859442f9ed1278bbdade68490a Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 26 May 2026 12:46:41 +0100 Subject: [PATCH 21/31] Use AutoMooncake in constrained.md's logdensity_and_gradient `AbstractPPL.prepare` only ships an `AutoMooncake` extension, so the previous `AutoForwardDiff()` call threw `MethodError` during `FisherMinBatchMatch`. Switch the import and the `prepare` call to `Mooncake`; the `Bijectors.Stacked` overlay defined in basic.md's `@setup` is already in effect by the time constrained.md runs. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/src/tutorials/constrained.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/src/tutorials/constrained.md b/docs/src/tutorials/constrained.md index 73a060902..6c5365c96 100644 --- a/docs/src/tutorials/constrained.md +++ b/docs/src/tutorials/constrained.md @@ -149,7 +149,7 @@ The rest can be done by simply copy-pasting the code below: ```@example constraints using ADTypes using AbstractPPL: AbstractPPL -using ForwardDiff: ForwardDiff +using Mooncake: Mooncake struct TransformedLogDensityProblem{Prob,BInv} prob::Prob @@ -172,8 +172,9 @@ function LogDensityProblems.logdensity_and_gradient( prob_trans::TransformedLogDensityProblem, θ_trans ) f = Base.Fix1(LogDensityProblems.logdensity, prob_trans) - prep = AbstractPPL.prepare(AutoForwardDiff(), f, θ_trans) - return AbstractPPL.value_and_gradient!!(prep, θ_trans) + x = collect(θ_trans) # AutoMooncake requires a dense vector + prep = AbstractPPL.prepare(AutoMooncake(), f, x) + return AbstractPPL.value_and_gradient!!(prep, x) end function LogDensityProblems.dimension(prob_trans::TransformedLogDensityProblem) From f7f89117d1fea96dc6a69d65f5611e6b3c83b041 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 26 May 2026 15:22:59 +0100 Subject: [PATCH 22/31] formatting --- docs/src/klminrepgraddescent.md | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/docs/src/klminrepgraddescent.md b/docs/src/klminrepgraddescent.md index 7fd22abbc..1c75e55f1 100644 --- a/docs/src/klminrepgraddescent.md +++ b/docs/src/klminrepgraddescent.md @@ -150,10 +150,7 @@ Recall that the original ADVI objective with a closed-form entropy (CFE) is give n_montecarlo = 16; cfe = KLMinRepGradDescent( - AutoMooncake(); - entropy=ClosedFormEntropy(), - optimizer=Adam(1e-2), - operator=ClipScale(), + AutoMooncake(); entropy=ClosedFormEntropy(), optimizer=Adam(1e-2), operator=ClipScale() ) nothing ``` @@ -264,7 +261,7 @@ using StatsFuns # QMC samples are inputs to AD, not parameters; declare them non-differentiable # so Mooncake doesn't trace through Sobol/Owen bit-twiddling. -Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(QuasiMonteCarlo.sample), Vararg} +Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(QuasiMonteCarlo.sample),Vararg} qmcrng = SobolSample(; R=OwenScramble(; base=2, pad=32)) From b4b2b47ea6f071ed7be9e4bb8a3afc467b4addc7 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 28 May 2026 19:47:58 +0100 Subject: [PATCH 23/31] Address Xianda's review on DynamicPPL ext + restore ReverseDiff benchmarks - ext: switch from deprecated `evaluate!!` to `DynamicPPL.logdensity_internal`, using a `WeightedLogJoint` getlogdensity that reads `loglikeadj` through the closure's Ref so subsample updates are observed without rebuilding AD prep. - ext: store the original (unsubsampled) model so repeated `subsample` calls compute `n_datapoints` from the full dataset, not the previous batch. - ext: broaden `capabilities` where-clause to dispatch on `adtype=nothing`. - ext: drop redundant `varinfo`/`accs` fields; store `dim::Int` directly. - docs/basic.md: dedent the FullRankGaussian intro so it is body text, not a `[^KW2014]` footnote continuation. - bench: re-add `AutoReverseDiff(compile=true)` to the AD list; pull `DifferentiationInterface` into bench deps so non-Mooncake, non-ForwardDiff backends find `AbstractPPL.prepare`. Co-Authored-By: Claude Opus 4.7 (1M context) --- bench/Project.toml | 4 ++ bench/benchmarks.jl | 4 +- docs/src/tutorials/basic.md | 3 +- ext/AdvancedVIDynamicPPLExt.jl | 95 +++++++++++++++++++++++++--------- src/AdvancedVI.jl | 26 +++++----- 5 files changed, 92 insertions(+), 40 deletions(-) diff --git a/bench/Project.toml b/bench/Project.toml index 1d3e954a1..a5fc803e1 100644 --- a/bench/Project.toml +++ b/bench/Project.toml @@ -2,6 +2,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" @@ -12,6 +13,7 @@ LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" @@ -19,6 +21,7 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" ADTypes = "1" AdvancedVI = "0.7, 0.6" BenchmarkTools = "1" +DifferentiationInterface = "0.6, 0.7" Distributions = "0.25.111" DistributionsAD = "0.6" Enzyme = "0.13.7" @@ -29,6 +32,7 @@ LogDensityProblems = "2" Mooncake = "0.4.5, 0.5" Optimisers = "0.3, 0.4" Random = "1" +ReverseDiff = "1" SimpleUnPack = "1" StableRNGs = "1" julia = "1.10, 1.11.2" diff --git a/bench/benchmarks.jl b/bench/benchmarks.jl index 878268faf..241a14081 100644 --- a/bench/benchmarks.jl +++ b/bench/benchmarks.jl @@ -3,7 +3,8 @@ using AdvancedVI using BenchmarkTools using Distributions using DistributionsAD -using Enzyme, ForwardDiff, Mooncake +using DifferentiationInterface # provides AbstractPPL.prepare for non-Mooncake, non-ForwardDiff backends +using Enzyme, ForwardDiff, Mooncake, ReverseDiff using FillArrays using InteractiveUtils using LinearAlgebra @@ -69,6 +70,7 @@ begin ], (adname, adtype) in [ ("Mooncake", AutoMooncake()), + ("ReverseDiff", AutoReverseDiff(; compile=true)), # ("Enzyme", AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)), ], (familyname, family) in [ diff --git a/docs/src/tutorials/basic.md b/docs/src/tutorials/basic.md index ecf278289..11ad44758 100644 --- a/docs/src/tutorials/basic.md +++ b/docs/src/tutorials/basic.md @@ -215,7 +215,8 @@ If the problem provides only a zeroth-order [capability](https://www.tamaspapp.e [^TL2014]: Titsias, M., & Lázaro-Gredilla, M. (2014, June). Doubly stochastic variational Bayes for non-conjugate inference. In *International Conference on Machine Learning*. PMLR. [^RMW2014]: Rezende, D. J., Mohamed, S., & Wierstra, D. (2014, June). Stochastic backpropagation and approximate inference in deep generative models. In *International Conference on Machine Learning*. PMLR. [^KW2014]: Kingma, D. P., & Welling, M. (2014). Auto-encoding variational bayes. In *International Conference on Learning Representations*. - For the variational family, we will consider a `FullRankGaussian` approximation: +For the variational family, we will consider a `FullRankGaussian` approximation: + ```@example basic using LinearAlgebra diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index 4363f2e3d..e826bbed8 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -13,17 +13,25 @@ function adtype_capabilities(::Type{<:ADTypes.AbstractADType}) return LogDensityProblems.LogDensityOrder{1}() end -function logdensity_impl( - params, model::DynamicPPL.Model, loglikeadj::Real, varinfo::DynamicPPL.AbstractVarInfo -) - vi = DynamicPPL.unflatten!!(varinfo, params) - _, vi = DynamicPPL.evaluate!!(model, vi) +# `getlogdensity` callable for `DynamicPPL.logdensity_internal`: reads the +# current `loglikeadj` through a Ref so the mutation done by `subsample` is +# observed without rebuilding any AD prep. +struct WeightedLogJoint{R} + loglikeadj_ref::R +end +function (g::WeightedLogJoint)(vi) loglike = DynamicPPL.getloglikelihood(vi) logprior = DynamicPPL.getlogprior(vi) logjac = DynamicPPL.getlogjac(vi) - return convert(eltype(params), loglikeadj) * loglike + logprior - logjac + return g.loglikeadj_ref[] * loglike + logprior - logjac end +const _DEFAULT_LDF_ACCS = DynamicPPL.AccumulatorTuple(( + DynamicPPL.LogPriorAccumulator(), + DynamicPPL.LogJacobianAccumulator(), + DynamicPPL.LogLikelihoodAccumulator(), +)) + function subsample_dynamicpplmodel( model::DynamicPPL.Model{F,A,D,M,Ta,Td,Ctx,Threaded}, batch ) where {F,A,D,M,Ta,Td,Ctx,Threaded} @@ -31,28 +39,38 @@ function subsample_dynamicpplmodel( return DynamicPPL.Model{Threaded}(model.f, model.args, new_kwargs, model.context) end -# `model_ref`/`loglikeadj_ref` are mutated in place by `subsample`; the closure -# inside `prep_grad`/`prep_hess` reads through them so the prep stays valid -# across subsampling steps. +# `model` is the original (unsubsampled) source of truth for `subsample` — it +# needs the full-dataset length on every call. `model_ref`/`loglikeadj_ref` are +# mutated in place by `subsample`; the closure inside `prep_grad`/`prep_hess` +# reads through them so the prep stays valid across subsampling steps. # # `model_ref` is typed `Ref{Any}` because `subsample_dynamicpplmodel` returns a # `DynamicPPL.Model` whose `defaults` NamedTuple type varies with the batch — a -# typed `Ref{<:DynamicPPL.Model}` would throw on reassignment. The tradeoff is a -# dynamic dispatch on each `prob.model_ref[]` read; do not "tighten" it. +# typed `Ref{<:DynamicPPL.Model}` would throw on reassignment. The tradeoff is +# a dynamic dispatch on each `prob.model_ref[]` read; do not "tighten" it. +# Compiled-tape backends (e.g. `AutoReverseDiff(compile=true)`) bake the deref +# into the tape at prep time and won't see `subsample` updates — they need a +# fresh `prob` per batch change. struct DynamicPPLModelLogDensityFunction{ Model<:DynamicPPL.Model, LogLikeAdj<:Real, - VarInfo<:DynamicPPL.AbstractVarInfo, + Ranges<:DynamicPPL.VarNamedTuple, + Strategy<:DynamicPPL.AbstractTransformStrategy, + GetLogDensity, ADType<:Union{Nothing,ADTypes.AbstractADType}, PrepGrad, PrepHess, } + model::Model model_ref::Ref{Any} loglikeadj_ref::Ref{LogLikeAdj} - varinfo::VarInfo + ranges_and_transforms::Ranges + transform_strategy::Strategy + getlogdensity::GetLogDensity adtype::ADType prep_grad::PrepGrad prep_hess::PrepHess + dim::Int end function DynamicPPLModelLogDensityFunction( @@ -72,11 +90,24 @@ function DynamicPPLModelLogDensityFunction( subsample_dynamicpplmodel(model, batch) end - params = collect(varinfo[:]) + ranges_and_transforms, params = DynamicPPL.get_rat_and_samplevec(varinfo.values) + transform_strategy = DynamicPPL.infer_transform_strategy_from_values( + ranges_and_transforms + ) + model_ref = Ref{Any}(model_sub) adj0 = float(loglikeadj) loglikeadj_ref = Ref(adj0) - f = params -> logdensity_impl(params, model_ref[], loglikeadj_ref[], varinfo) + getlogdensity = WeightedLogJoint(loglikeadj_ref) + f = + params -> DynamicPPL.logdensity_internal( + params, + model_ref[], + getlogdensity, + ranges_and_transforms, + transform_strategy, + _DEFAULT_LDF_ACCS, + ) cap = adtype_capabilities(typeof(adtype)) prep_grad = if cap >= LogDensityProblems.LogDensityOrder{1}() @@ -98,17 +129,35 @@ function DynamicPPLModelLogDensityFunction( return DynamicPPLModelLogDensityFunction{ typeof(model), typeof(adj0), - typeof(varinfo), + typeof(ranges_and_transforms), + typeof(transform_strategy), + typeof(getlogdensity), typeof(adtype), typeof(prep_grad), typeof(prep_hess), }( - model_ref, loglikeadj_ref, varinfo, adtype, prep_grad, prep_hess + model, + model_ref, + loglikeadj_ref, + ranges_and_transforms, + transform_strategy, + getlogdensity, + adtype, + prep_grad, + prep_hess, + length(params), ) end function LogDensityProblems.logdensity(prob::DynamicPPLModelLogDensityFunction, params) - return logdensity_impl(params, prob.model_ref[], prob.loglikeadj_ref[], prob.varinfo) + return DynamicPPL.logdensity_internal( + params, + prob.model_ref[], + prob.getlogdensity, + prob.ranges_and_transforms, + prob.transform_strategy, + _DEFAULT_LDF_ACCS, + ) end # `!!` may alias internal buffers of `prep_*`; copy so callers can retain the @@ -128,8 +177,8 @@ function LogDensityProblems.logdensity_gradient_and_hessian( end function LogDensityProblems.capabilities( - ::Type{<:DynamicPPLModelLogDensityFunction{M,L,V,ADType,PG,PH}} -) where {M,L,V,ADType<:ADTypes.AbstractADType,PG,PH} + ::Type{<:DynamicPPLModelLogDensityFunction{Model,L,R,S,G,ADType,PG,PH}} +) where {Model,L,R,S,G,ADType<:Union{Nothing,ADTypes.AbstractADType},PG,PH} return if PH !== Nothing LogDensityProblems.LogDensityOrder{2}() elseif PG !== Nothing @@ -139,12 +188,10 @@ function LogDensityProblems.capabilities( end end -function LogDensityProblems.dimension(prob::DynamicPPLModelLogDensityFunction) - return length(prob.varinfo[:]) -end +LogDensityProblems.dimension(prob::DynamicPPLModelLogDensityFunction) = prob.dim function AdvancedVI.subsample(prob::DynamicPPLModelLogDensityFunction, batch) - model = prob.model_ref[] + model = prob.model # full dataset; do not read from `model_ref[]` here if !haskey(model.defaults, :datapoints) throw( diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 568c3189d..afbaec191 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -45,7 +45,7 @@ Evaluate the value and gradient of a function `f` at `x` using the automatic dif - `f`: Function subject to differentiation. - `x`: The point to evaluate the gradient. - `aux`: Auxiliary input passed to `f`. -- `prep`: Output of `_prepare_gradient`. +- `prep`: Output of `_prepare_gradient` (`nothing` means "re-prepare each call"). - `out::DiffResults.MutableDiffResult`: Buffer to contain the output gradient and function value. """ function _value_and_gradient!( @@ -58,9 +58,17 @@ function _value_and_gradient!( return out end +function _value_and_gradient!( + f, out::DiffResults.MutableDiffResult, ::Nothing, ad::ADTypes.AbstractADType, x, aux +) + return _value_and_gradient!(f, out, ad, x, aux) +end + function _value_and_gradient!( f, out::DiffResults.MutableDiffResult, prep, ad::ADTypes.AbstractADType, x, aux ) + # `context[1]` is the `Ref(aux)` placed there by `_prepare_gradient`; mutating + # it lets a cached prep see the new `aux` without rebuilding AD machinery. prep.evaluator.context[1][] = aux val, grad = AbstractPPL.value_and_gradient!!(prep, x) DiffResults.value!(out, val) @@ -68,19 +76,6 @@ function _value_and_gradient!( return out end -# Compiled-tape ReverseDiff bakes `aref[]` into the tape at prep time, so a -# cached prep would feed stale `aux` to AD. Re-prepare each call instead. -function _value_and_gradient!( - f, - out::DiffResults.MutableDiffResult, - ::Nothing, - ad::ADTypes.AutoReverseDiff{true}, - x, - aux, -) - return _value_and_gradient!(f, out, ad, x, aux) -end - """ _prepare_gradient(f, ad, x, aux) @@ -96,6 +91,9 @@ function _prepare_gradient(f, ad::ADTypes.AbstractADType, x, aux) return AbstractPPL.prepare(ad, (x, aref) -> f(x, aref[]), x; context=(Ref(aux),)) end +# Compiled-tape ReverseDiff bakes context values into the tape at prep time, +# so the `Ref(aux)` trick above would feed stale `aux` to AD. Skip the cache +# and let `_value_and_gradient!` re-prepare on every call. _prepare_gradient(::Any, ::ADTypes.AutoReverseDiff{true}, ::Any, ::Any) = nothing """ From 4dab3e7e85e533e5301ae4f5ec7634eebe067ec9 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 28 May 2026 19:55:05 +0100 Subject: [PATCH 24/31] Tighten DynamicPPL ext design-rationale comments Trim the struct-level WHY block and the inline `prob.model` comment to keep just the rationale (full-dataset invariant, `Ref{Any}` choice, compiled-tape caveat) without restating control flow. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AdvancedVIDynamicPPLExt.jl | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index e826bbed8..816ceac69 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -39,18 +39,15 @@ function subsample_dynamicpplmodel( return DynamicPPL.Model{Threaded}(model.f, model.args, new_kwargs, model.context) end -# `model` is the original (unsubsampled) source of truth for `subsample` — it -# needs the full-dataset length on every call. `model_ref`/`loglikeadj_ref` are -# mutated in place by `subsample`; the closure inside `prep_grad`/`prep_hess` -# reads through them so the prep stays valid across subsampling steps. +# `model` is the original (unsubsampled) source of truth; `subsample` must read +# it (not `model_ref[]`) to get the full-dataset length on every call. +# `model_ref`/`loglikeadj_ref` are mutated in place by `subsample` so the +# closure inside `prep_grad`/`prep_hess` stays valid across subsampling steps. # -# `model_ref` is typed `Ref{Any}` because `subsample_dynamicpplmodel` returns a -# `DynamicPPL.Model` whose `defaults` NamedTuple type varies with the batch — a -# typed `Ref{<:DynamicPPL.Model}` would throw on reassignment. The tradeoff is -# a dynamic dispatch on each `prob.model_ref[]` read; do not "tighten" it. -# Compiled-tape backends (e.g. `AutoReverseDiff(compile=true)`) bake the deref -# into the tape at prep time and won't see `subsample` updates — they need a -# fresh `prob` per batch change. +# `model_ref` is `Ref{Any}` because `subsample_dynamicpplmodel`'s output type +# varies with the batch; a typed Ref would throw on reassignment. Compiled-tape +# backends (e.g. `AutoReverseDiff(compile=true)`) bake the deref into the tape +# at prep time and won't observe `subsample` updates without a fresh `prob`. struct DynamicPPLModelLogDensityFunction{ Model<:DynamicPPL.Model, LogLikeAdj<:Real, @@ -191,7 +188,7 @@ end LogDensityProblems.dimension(prob::DynamicPPLModelLogDensityFunction) = prob.dim function AdvancedVI.subsample(prob::DynamicPPLModelLogDensityFunction, batch) - model = prob.model # full dataset; do not read from `model_ref[]` here + model = prob.model # full dataset — `model_ref[]` would already be subsampled if !haskey(model.defaults, :datapoints) throw( From daf504e1a9ce9e92b7b7591e202d0b35b3535717 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 28 May 2026 19:57:53 +0100 Subject: [PATCH 25/31] Move basic.md footnotes to end of file JuliaFormatter v1 with `format_markdown=true` re-indented the `FullRankGaussian` intro line into the `[^KW2014]` footnote because the footnotes were sandwiched between body text. Moving them to the end of the document removes the ambiguity, so both Xianda's review fix and the formatter check are satisfied. Co-Authored-By: Claude Opus 4.7 (1M context) --- bench/benchmarks.jl | 2 +- docs/src/tutorials/basic.md | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/bench/benchmarks.jl b/bench/benchmarks.jl index 241a14081..4dd9d4a30 100644 --- a/bench/benchmarks.jl +++ b/bench/benchmarks.jl @@ -70,7 +70,7 @@ begin ], (adname, adtype) in [ ("Mooncake", AutoMooncake()), - ("ReverseDiff", AutoReverseDiff(; compile=true)), + ("ReverseDiff", AutoReverseDiff()), # ("Enzyme", AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)), ], (familyname, family) in [ diff --git a/docs/src/tutorials/basic.md b/docs/src/tutorials/basic.md index 11ad44758..0c35d6351 100644 --- a/docs/src/tutorials/basic.md +++ b/docs/src/tutorials/basic.md @@ -212,9 +212,6 @@ Here, the projection operator `ClipScale` ensures this. `KLMinRepGradDescent`, in particular, assumes that the target `LogDensityProblem` is differentiable. If the problem provides only a zeroth-order [capability](https://www.tamaspapp.eu/LogDensityProblems.jl/dev/#LogDensityProblems.capabilities), `AdvancedVI` will differentiate through `LogDensityProblems.logdensity` directly using the AD backend supplied to the algorithm. -[^TL2014]: Titsias, M., & Lázaro-Gredilla, M. (2014, June). Doubly stochastic variational Bayes for non-conjugate inference. In *International Conference on Machine Learning*. PMLR. -[^RMW2014]: Rezende, D. J., Mohamed, S., & Wierstra, D. (2014, June). Stochastic backpropagation and approximate inference in deep generative models. In *International Conference on Machine Learning*. PMLR. -[^KW2014]: Kingma, D. P., & Welling, M. (2014). Auto-encoding variational bayes. In *International Conference on Learning Representations*. For the variational family, we will consider a `FullRankGaussian` approximation: ```@example basic @@ -373,3 +370,7 @@ nothing ![](basic_example_acc.svg) Clearly, the accuracy is improving over time. + +[^TL2014]: Titsias, M., & Lázaro-Gredilla, M. (2014, June). Doubly stochastic variational Bayes for non-conjugate inference. In *International Conference on Machine Learning*. PMLR. +[^RMW2014]: Rezende, D. J., Mohamed, S., & Wierstra, D. (2014, June). Stochastic backpropagation and approximate inference in deep generative models. In *International Conference on Machine Learning*. PMLR. +[^KW2014]: Kingma, D. P., & Welling, M. (2014). Auto-encoding variational bayes. In *International Conference on Learning Representations*. From 1a982d3d3f3c2cf9895875bf6f232ae56f69787e Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 1 Jun 2026 15:50:26 +0100 Subject: [PATCH 26/31] Address remaining review feedback from sunxd3 - src/AdvancedVI.jl: `AutoReverseDiff(; compile=true)` now throws an `ArgumentError` from `_prepare_gradient` rather than silently returning `nothing` and re-preparing on every call. Compiled tapes freeze captured values at preparation time, so subsequent iterations would differentiate against stale data and silently produce incorrect gradients. Drops the now-dead `_value_and_gradient!(::Nothing, ...)` overload and the `(nothing means re-prepare)` mention in its docstring. - src/families/location_scale_low_rank.jl: clarify that the `Symmetric`-instead-of-`Hermitian` workaround exists because Mooncake lacks an `rrule` for the `Hermitian` path. - test/Project.toml: drop redundant `AdvancedVI` entry from `[deps]`; `Pkg.test` injects the package under test automatically. --- src/AdvancedVI.jl | 28 ++++++++++++++++--------- src/families/location_scale_low_rank.jl | 2 +- test/Project.toml | 1 - 3 files changed, 19 insertions(+), 12 deletions(-) diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index afbaec191..323278066 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -45,7 +45,7 @@ Evaluate the value and gradient of a function `f` at `x` using the automatic dif - `f`: Function subject to differentiation. - `x`: The point to evaluate the gradient. - `aux`: Auxiliary input passed to `f`. -- `prep`: Output of `_prepare_gradient` (`nothing` means "re-prepare each call"). +- `prep`: Output of `_prepare_gradient`. - `out::DiffResults.MutableDiffResult`: Buffer to contain the output gradient and function value. """ function _value_and_gradient!( @@ -58,12 +58,6 @@ function _value_and_gradient!( return out end -function _value_and_gradient!( - f, out::DiffResults.MutableDiffResult, ::Nothing, ad::ADTypes.AbstractADType, x, aux -) - return _value_and_gradient!(f, out, ad, x, aux) -end - function _value_and_gradient!( f, out::DiffResults.MutableDiffResult, prep, ad::ADTypes.AbstractADType, x, aux ) @@ -92,9 +86,20 @@ function _prepare_gradient(f, ad::ADTypes.AbstractADType, x, aux) end # Compiled-tape ReverseDiff bakes context values into the tape at prep time, -# so the `Ref(aux)` trick above would feed stale `aux` to AD. Skip the cache -# and let `_value_and_gradient!` re-prepare on every call. -_prepare_gradient(::Any, ::ADTypes.AutoReverseDiff{true}, ::Any, ::Any) = nothing +# so the `Ref(aux)` trick above would feed stale `aux` to AD. Reject it +# loudly rather than silently producing wrong gradients. +function _prepare_gradient(::Any, ::ADTypes.AutoReverseDiff{true}, ::Any, ::Any) + throw( + ArgumentError( + "`AutoReverseDiff(; compile=true)` is not safe to use with AdvancedVI: " * + "compiled tapes freeze captured values at preparation time, so " * + "subsequent iterations differentiate against stale data and silently " * + "produce incorrect gradients. Use `AutoReverseDiff(; compile=false)` " * + "(or another reverse-mode backend such as `AutoMooncake` or " * + "`AutoEnzyme`) instead.", + ), + ) +end """ restructure_ad_forward(adtype, restructure, params) @@ -305,6 +310,9 @@ export estimate_objective Inform `model` or `q` to only use the data points designated by the iterable collection `batch`. For `model`, the log-density should also be adjusted to account for the change in number of data points. + +Implementations may mutate `model`/`q` in place and return the same object; +callers should not retain the pre-call binding. """ subsample(model_or_q::Any, ::Any) = model_or_q diff --git a/src/families/location_scale_low_rank.jl b/src/families/location_scale_low_rank.jl index 6703bcea3..3414a5e45 100644 --- a/src/families/location_scale_low_rank.jl +++ b/src/families/location_scale_low_rank.jl @@ -36,7 +36,7 @@ function StatsBase.entropy(q::MvLocationScaleLowRank) (; location, scale_diag, scale_factors, dist) = q n_dims = length(location) scale_diag2 = scale_diag .* scale_diag - # `Symmetric` rather than `Hermitian`: Mooncake lacks a rule for the `Hermitian` path. + # `Symmetric` rather than `Hermitian`: Mooncake lacks an `rrule` for the `Hermitian` path. UtDinvU = Symmetric(scale_factors' * (scale_factors ./ scale_diag2)) logdetΣ = 2 * sum(log.(scale_diag)) + logdet(I + UtDinvU) return n_dims * convert(eltype(location), entropy(dist)) + logdetΣ / 2 diff --git a/test/Project.toml b/test/Project.toml index 55acfbe31..cab7aff1e 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,7 +1,6 @@ [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" -AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" From b911715cfe967bbd0c23f1fa298be39656b7d4be Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 1 Jun 2026 16:19:47 +0100 Subject: [PATCH 27/31] Revert benchmark workflow back to LTS Julia MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The PR previously bumped Benchmark.yml to Julia 1 because of a DynamicPPL main `[sources]` pin that no longer exists (the PR now tracks the released DynamicPPL 0.42). Main's benchmark baseline still runs on LTS, so the PR running on 1.12 produced a false-positive 2× regression on ReverseDiff (which is slower on 1.12). Mooncake numbers in the same comparison actually improved, confirming the regression was a Julia-version artifact, not a code change. --- .github/workflows/Benchmark.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/Benchmark.yml b/.github/workflows/Benchmark.yml index 67e53985c..14e7d57b6 100644 --- a/.github/workflows/Benchmark.yml +++ b/.github/workflows/Benchmark.yml @@ -22,7 +22,7 @@ jobs: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v2 with: - version: '1' + version: 'lts' arch: x64 - uses: actions/cache@v4 env: From bd6160fd79b974318005b31efd41938d69670afd Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 3 Jun 2026 12:20:28 +0100 Subject: [PATCH 28/31] Drop unused test deps and clarify subsample docstring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `AbstractPPL` and `Bijectors` were added to test/Project.toml when this PR carried `[sources]` pins against dev branches. The pins are gone; neither package is referenced from test code, and DynamicPPL 0.42 already pins `Bijectors = "0.16"` transitively, so the test-side declarations are dead weight. Also reword `subsample`'s docstring: the prior wording told callers not to "retain the pre-call binding," but when an implementation mutates and returns the same object the binding is still valid — it just points to a now-mutated object. The intent is that the returned value may alias the input. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/AdvancedVI.jl | 4 ++-- test/Project.toml | 3 --- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 323278066..03c079ac6 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -311,8 +311,8 @@ export estimate_objective Inform `model` or `q` to only use the data points designated by the iterable collection `batch`. For `model`, the log-density should also be adjusted to account for the change in number of data points. -Implementations may mutate `model`/`q` in place and return the same object; -callers should not retain the pre-call binding. +Implementations may mutate `model`/`q` in place and return the same object, +so callers must not assume the returned value is distinct from the input. """ subsample(model_or_q::Any, ::Any) = model_or_q diff --git a/test/Project.toml b/test/Project.toml index cab7aff1e..7cb2e8a7a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,7 +1,5 @@ [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" -AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" -Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" @@ -26,7 +24,6 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" [compat] ADTypes = "0.2.1, 1" -Bijectors = "0.16" DiffResults = "1" DifferentiationInterface = "0.6, 0.7" Distributions = "0.25.111" From 428d5c055c825376338d19636896fb6546a24e7b Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 3 Jun 2026 17:40:15 +0100 Subject: [PATCH 29/31] Address review feedback: hoist prep, drop Mooncake overlay, tighten ext constrained.md: hoist AbstractPPL.prepare into TransformedLogDensityProblem's constructor (Red-Portal #3348649965) so the AD backend is traced once at setup rather than rebuilt per logdensity_and_gradient call. basic.md: drop the @mooncake_overlay workaround and its
explainer. docs/Project.toml [sources] pins Mooncake to chalk-lab/Mooncake.jl's fix-1191-fcodual-unionall branch, which extends the @isdefined(P) guard down to the generic codual_type/fcodual_type methods and so resolves the CategoricalArrays unbound-TypeVar interaction at source. ext/src cleanups (from scrutinise pass): tighten WeightedLogJoint's TypeVar bound to Base.RefValue{<:Real}, collapse the intermediate adj0 in the ext constructor, and replace the brittle AD-backend enumeration in _value_and_gradient!'s docstring with a generic AbstractPPL.prepare reference. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/Project.toml | 3 ++ docs/src/tutorials/basic.md | 51 ------------------------------- docs/src/tutorials/constrained.md | 28 +++++++++++------ ext/AdvancedVIDynamicPPLExt.jl | 14 ++++----- src/AdvancedVI.jl | 14 +++------ 5 files changed, 32 insertions(+), 78 deletions(-) diff --git a/docs/Project.toml b/docs/Project.toml index ee4daaa3f..29180fd7d 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -21,6 +21,9 @@ QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b" StanLogDensityProblems = "a545de4d-8dba-46db-9d34-4e41d3f07807" StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c" +[sources] +Mooncake = {url = "https://github.com/chalk-lab/Mooncake.jl", rev = "fix-1191-fcodual-unionall"} + [compat] ADTypes = "1" AbstractPPL = "0.15" diff --git a/docs/src/tutorials/basic.md b/docs/src/tutorials/basic.md index 0c35d6351..5fb315918 100644 --- a/docs/src/tutorials/basic.md +++ b/docs/src/tutorials/basic.md @@ -120,57 +120,6 @@ y = Vector{Bool}(data[:, end] .== "Mine") nothing ``` -```@setup basic -using Bijectors: Bijectors -using Mooncake: Mooncake - -Mooncake.@mooncake_overlay function Bijectors._with_logabsdet_jacobian( - sb::Bijectors.Stacked, x::AbstractVector -) - T = eltype(x) - y = T[] - logjac = zero(T) - for i in eachindex(sb.bs) - yi, lji = Bijectors.with_logabsdet_jacobian(sb.bs[i], x[sb.ranges_in[i]]) - y = vcat(y, yi) - logjac = logjac + sum(lji) - end - return (y, logjac) -end -``` - -```@raw html -
-Why does Mooncake need help differentiating Bijectors.Stacked when OpenML is loaded? (click to expand) -``` - -Loading `OpenML` transitively brings in `CategoricalArrays`, which adds -`Base.similar` methods with unbound type parameters. Those break Mooncake's -compilation of `collect(::Base.Generator)` — used by `Bijectors.Stacked`'s -`_with_logabsdet_jacobian` via `map(zip(...)) do ...`. The setup block above -sidesteps it with `Mooncake.@mooncake_overlay`, replacing the implementation -with an explicit loop: - -```julia -Mooncake.@mooncake_overlay function Bijectors._with_logabsdet_jacobian( - sb::Bijectors.Stacked, x::AbstractVector -) - T = eltype(x) - y = T[] - logjac = zero(T) - for i in eachindex(sb.bs) - yi, lji = Bijectors.with_logabsdet_jacobian(sb.bs[i], x[sb.ranges_in[i]]) - y = vcat(y, yi) - logjac = logjac + sum(lji) - end - return (y, logjac) -end -``` - -```@raw html -
-``` - Let's apply some basic pre-processing and add an intercept column: ```@example basic diff --git a/docs/src/tutorials/constrained.md b/docs/src/tutorials/constrained.md index 6c5365c96..ca7462c59 100644 --- a/docs/src/tutorials/constrained.md +++ b/docs/src/tutorials/constrained.md @@ -151,30 +151,37 @@ using ADTypes using AbstractPPL: AbstractPPL using Mooncake: Mooncake -struct TransformedLogDensityProblem{Prob,BInv} +struct TransformedLogDensityProblem{Prob,BInv,Prep} prob::Prob binv::BInv + prep::Prep end -function TransformedLogDensityProblem(prob) +function TransformedLogDensityProblem(prob; adtype::ADTypes.AbstractADType=AutoMooncake()) b = Bijectors.bijector(prob) binv = Bijectors.inverse(b) - return TransformedLogDensityProblem{typeof(prob),typeof(binv)}(prob, binv) + d_unc = prod(Bijectors.output_size(b, (LogDensityProblems.dimension(prob),))) + + f = let prob = prob, binv = binv + function (θ_trans) + θ, logabsdetjac = Bijectors.with_logabsdet_jacobian(binv, θ_trans) + return LogDensityProblems.logdensity(prob, θ) + logabsdetjac + end + end + prep = AbstractPPL.prepare(adtype, f, zeros(d_unc)) + return TransformedLogDensityProblem(prob, binv, prep) end function LogDensityProblems.logdensity(prob_trans::TransformedLogDensityProblem, θ_trans) - (; prob, binv) = prob_trans - θ, logabsdetjac = Bijectors.with_logabsdet_jacobian(binv, θ_trans) - return LogDensityProblems.logdensity(prob, θ) + logabsdetjac + return prob_trans.prep(θ_trans) end function LogDensityProblems.logdensity_and_gradient( prob_trans::TransformedLogDensityProblem, θ_trans ) - f = Base.Fix1(LogDensityProblems.logdensity, prob_trans) - x = collect(θ_trans) # AutoMooncake requires a dense vector - prep = AbstractPPL.prepare(AutoMooncake(), f, x) - return AbstractPPL.value_and_gradient!!(prep, x) + # `prep` was built for a dense `Vector`; collect view-typed inputs + # (e.g. minibatch column views) to match the AD backend's expectations. + return AbstractPPL.value_and_gradient!!(prob_trans.prep, collect(θ_trans)) end function LogDensityProblems.dimension(prob_trans::TransformedLogDensityProblem) @@ -191,6 +198,7 @@ nothing ``` Wrapping `prob` with `TransformedLogDensityProblem` yields our unconstrained posterior. +The `AbstractPPL.prepare` call traces the AD backend once at construction time, so the per-iteration `logdensity_and_gradient` call just runs the cached evaluator. ```@example constraints prob_trans = TransformedLogDensityProblem(prob) diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index 816ceac69..445dc90ae 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -16,7 +16,7 @@ end # `getlogdensity` callable for `DynamicPPL.logdensity_internal`: reads the # current `loglikeadj` through a Ref so the mutation done by `subsample` is # observed without rebuilding any AD prep. -struct WeightedLogJoint{R} +struct WeightedLogJoint{R<:Base.RefValue{<:Real}} loglikeadj_ref::R end function (g::WeightedLogJoint)(vi) @@ -43,11 +43,10 @@ end # it (not `model_ref[]`) to get the full-dataset length on every call. # `model_ref`/`loglikeadj_ref` are mutated in place by `subsample` so the # closure inside `prep_grad`/`prep_hess` stays valid across subsampling steps. -# # `model_ref` is `Ref{Any}` because `subsample_dynamicpplmodel`'s output type -# varies with the batch; a typed Ref would throw on reassignment. Compiled-tape -# backends (e.g. `AutoReverseDiff(compile=true)`) bake the deref into the tape -# at prep time and won't observe `subsample` updates without a fresh `prob`. +# varies with the batch (a typed Ref would throw on reassignment), and because +# compiled-tape backends would otherwise bake the deref into the tape and miss +# the `subsample` update. struct DynamicPPLModelLogDensityFunction{ Model<:DynamicPPL.Model, LogLikeAdj<:Real, @@ -93,8 +92,7 @@ function DynamicPPLModelLogDensityFunction( ) model_ref = Ref{Any}(model_sub) - adj0 = float(loglikeadj) - loglikeadj_ref = Ref(adj0) + loglikeadj_ref = Ref(float(loglikeadj)) getlogdensity = WeightedLogJoint(loglikeadj_ref) f = params -> DynamicPPL.logdensity_internal( @@ -125,7 +123,7 @@ function DynamicPPLModelLogDensityFunction( end return DynamicPPLModelLogDensityFunction{ typeof(model), - typeof(adj0), + typeof(loglikeadj_ref[]), typeof(ranges_and_transforms), typeof(transform_strategy), typeof(getlogdensity), diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 03c079ac6..727fb2f5d 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -33,15 +33,11 @@ Evaluate the value and gradient of a function `f` at `x` using the automatic dif `f` may receive auxiliary input as `f(x,aux)`. # Arguments -- `ad::ADTypes.AbstractADType`: - automatic differentiation backend. Currently supports - `ADTypes.ForwardDiff()`, `ADTypes.ReverseDiff()`, - `ADTypes.AutoMooncake()` and - `ADTypes.AutoEnzyme(; - mode=Enzyme.set_runtime_activity(Enzyme.Reverse), - function_annotation=Enzyme.Const, - )`. - If one wants to use `AutoEnzyme`, please make sure to include the `set_runtime_activity` and `function_annotation` as shown above. +- `ad::ADTypes.AbstractADType`: automatic differentiation backend; any backend + with an `AbstractPPL.prepare` method is supported. `AutoEnzyme` requires + `set_runtime_activity` and `function_annotation=Enzyme.Const`. + `AutoReverseDiff(; compile=true)` is rejected because compiled tapes + freeze captured values at prep time. - `f`: Function subject to differentiation. - `x`: The point to evaluate the gradient. - `aux`: Auxiliary input passed to `f`. From 01c6b7d3af7b2a9198d274807eb00ad7bc7a5959 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 4 Jun 2026 13:10:16 +0100 Subject: [PATCH 30/31] Use eltype(loglikeadj_ref) for consistency Same file already uses `eltype(prob.loglikeadj_ref)` at line 204 to extract the Ref's element type; bring the constructor's type parameter computation in line with that idiom. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AdvancedVIDynamicPPLExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index 445dc90ae..687a36bb0 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -123,7 +123,7 @@ function DynamicPPLModelLogDensityFunction( end return DynamicPPLModelLogDensityFunction{ typeof(model), - typeof(loglikeadj_ref[]), + eltype(loglikeadj_ref), typeof(ranges_and_transforms), typeof(transform_strategy), typeof(getlogdensity), From ece02c196dc12ac282f7c11aef893df4851fc71c Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 4 Jun 2026 14:11:28 +0100 Subject: [PATCH 31/31] Require Mooncake 0.5.31, drop docs [sources] pin Mooncake 0.5.31 extends chalk-lab/Mooncake.jl#1192's `@isdefined(P)` guard to the generic `codual_type(::Type{P})` and `fcodual_type(::Type{P})` methods. That resolves the `UndefVarError: P not defined in static parameter matching` Mooncake threw when CategoricalArrays (loaded transitively via OpenML in basic.md) introduced loosely-constrained `similar` overloads into the matching set during AD rule derivation. With the fix in a released version, the docs no longer need the `[sources]` pin to the unmerged fork branch, and the docs build's `@mooncake_overlay` workaround on `Bijectors._with_logabsdet_jacobian` (removed in the previous commit) stays gone. Co-Authored-By: Claude Opus 4.7 (1M context) --- Project.toml | 2 +- docs/Project.toml | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index 7fe21d671..5999f0ee1 100644 --- a/Project.toml +++ b/Project.toml @@ -45,7 +45,7 @@ FillArrays = "1.3" Functors = "0.4, 0.5" LinearAlgebra = "1" LogDensityProblems = "2" -Mooncake = "0.4, 0.5" +Mooncake = "0.5.31" Optimisers = "0.2.16, 0.3, 0.4" ProgressMeter = "1.6" Random = "1" diff --git a/docs/Project.toml b/docs/Project.toml index 72ccef718..ddd05fdd8 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -21,9 +21,6 @@ QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b" StanLogDensityProblems = "a545de4d-8dba-46db-9d34-4e41d3f07807" StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c" -[sources] -Mooncake = {url = "https://github.com/chalk-lab/Mooncake.jl", rev = "fix-1191-fcodual-unionall"} - [compat] ADTypes = "1" AbstractPPL = "0.15" @@ -38,7 +35,7 @@ ForwardDiff = "0.10, 1" Functors = "0.5" JSON = "0.21, 1" LogDensityProblems = "2.1.1" -Mooncake = "0.4, 0.5" +Mooncake = "0.5.31" NormalizingFlows = "0.2.2" OpenML = "0.3" Optimisers = "0.3, 0.4"