First efforts towards DPPL 0.37 compat, WIP

Author: mhauru

Project.toml

@@ -64,7 +64,7 @@ Distributions = "0.25.77"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
 DynamicHMC = "3.4"
-DynamicPPL = "0.36.3"
+DynamicPPL = "0.37"
 EllipticalSliceSampling = "0.5, 1, 2"
 ForwardDiff = "0.10.3, 1"
 Libtask = "0.8.8"

ext/TuringOptimExt.jl

@@ -34,8 +34,8 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    ctx = Optimisation.OptimizationContext(DynamicPPL.LikelihoodContext())
-    f = Optimisation.OptimLogDensity(model, ctx)
+    vi = DynamicPPL.setaccs!!(VarInfo(model), (DynamicPPL.LogLikelihoodAccumulator(),))
+    f = Optimisation.OptimLogDensity(model, vi)
     init_vals = DynamicPPL.getparams(f.ldf)
     optimizer = Optim.LBFGS()
     return _mle_optimize(model, init_vals, optimizer, options; kwargs...)
@@ -57,8 +57,8 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    ctx = Optimisation.OptimizationContext(DynamicPPL.LikelihoodContext())
-    f = Optimisation.OptimLogDensity(model, ctx)
+    vi = DynamicPPL.setaccs!!(VarInfo(model), (DynamicPPL.LogLikelihoodAccumulator(),))
+    f = Optimisation.OptimLogDensity(model, vi)
     init_vals = DynamicPPL.getparams(f.ldf)
     return _mle_optimize(model, init_vals, optimizer, options; kwargs...)
 end
@@ -74,8 +74,9 @@ function Optim.optimize(
 end
 
 function _mle_optimize(model::DynamicPPL.Model, args...; kwargs...)
-    ctx = Optimisation.OptimizationContext(DynamicPPL.LikelihoodContext())
-    return _optimize(Optimisation.OptimLogDensity(model, ctx), args...; kwargs...)
+    vi = DynamicPPL.setaccs!!(VarInfo(model), (DynamicPPL.LogLikelihoodAccumulator(),))
+    f = Optimisation.OptimLogDensity(model, vi)
+    return _optimize(f, args...; kwargs...)
 end
 
 """
@@ -104,8 +105,8 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    ctx = Optimisation.OptimizationContext(DynamicPPL.DefaultContext())
-    f = Optimisation.OptimLogDensity(model, ctx)
+    vi = DynamicPPL.setaccs!!(VarInfo(model), (LogPriorWithoutJacobianAccumulator(), DynamicPPL.LogLikelihoodAccumulator(),))
+    f = Optimisation.OptimLogDensity(model, vi)
     init_vals = DynamicPPL.getparams(f.ldf)
     optimizer = Optim.LBFGS()
     return _map_optimize(model, init_vals, optimizer, options; kwargs...)
@@ -127,8 +128,8 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    ctx = Optimisation.OptimizationContext(DynamicPPL.DefaultContext())
-    f = Optimisation.OptimLogDensity(model, ctx)
+    vi = DynamicPPL.setaccs!!(VarInfo(model), (LogPriorWithoutJacobianAccumulator(), DynamicPPL.LogLikelihoodAccumulator(),))
+    f = Optimisation.OptimLogDensity(model, vi)
     init_vals = DynamicPPL.getparams(f.ldf)
     return _map_optimize(model, init_vals, optimizer, options; kwargs...)
 end
@@ -144,9 +145,11 @@ function Optim.optimize(
 end
 
 function _map_optimize(model::DynamicPPL.Model, args...; kwargs...)
-    ctx = Optimisation.OptimizationContext(DynamicPPL.DefaultContext())
-    return _optimize(Optimisation.OptimLogDensity(model, ctx), args...; kwargs...)
+    vi = DynamicPPL.setaccs!!(VarInfo(model), (LogPriorWithoutJacobianAccumulator(), DynamicPPL.LogLikelihoodAccumulator(),))
+    f = Optimisation.OptimLogDensity(model, vi)
+    return _optimize(f, args...; kwargs...)
 end
+
 """
     _optimize(f::OptimLogDensity, optimizer=Optim.LBFGS(), args...; kwargs...)
 
@@ -166,7 +169,7 @@ function _optimize(
     # whether initialisation is really necessary at all
     vi = DynamicPPL.unflatten(f.ldf.varinfo, init_vals)
     vi = DynamicPPL.link(vi, f.ldf.model)
-    f = Optimisation.OptimLogDensity(f.ldf.model, vi, f.ldf.context; adtype=f.ldf.adtype)
+    f = Optimisation.OptimLogDensity(f.ldf.model, vi; adtype=f.ldf.adtype)
     init_vals = DynamicPPL.getparams(f.ldf)
 
     # Optimize!
@@ -183,9 +186,7 @@ function _optimize(
     # Get the optimum in unconstrained space. `getparams` does the invlinking.
     vi = f.ldf.varinfo
     vi_optimum = DynamicPPL.unflatten(vi, M.minimizer)
-    logdensity_optimum = Optimisation.OptimLogDensity(
-        f.ldf.model, vi_optimum, f.ldf.context
-    )
+    logdensity_optimum = Optimisation.OptimLogDensity(f.ldf.model, vi_optimum; adtype=f.ldf.adtype)
     vns_vals_iter = Turing.Inference.getparams(f.ldf.model, vi_optimum)
     varnames = map(Symbol ∘ first, vns_vals_iter)
     vals = map(last, vns_vals_iter)

src/mcmc/Inference.jl

@@ -26,9 +26,6 @@ using DynamicPPL:
     SampleFromPrior,
     SampleFromUniform,
     DefaultContext,
-    PriorContext,
-    LikelihoodContext,
-    SamplingContext,
     set_flag!,
     unset_flag!
 using Distributions, Libtask, Bijectors

src/mcmc/ess.jl

@@ -49,7 +49,7 @@ function AbstractMCMC.step(
         rng,
         EllipticalSliceSampling.ESSModel(
             ESSPrior(model, spl, vi),
-            DynamicPPL.LogDensityFunction(
+            DynamicPPL.LogDensityFunction{:LogLikelihood}(
                 model, vi, DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext())
             ),
         ),
@@ -59,7 +59,7 @@ function AbstractMCMC.step(
 
     # update sample and log-likelihood
     vi = DynamicPPL.unflatten(vi, sample)
-    vi = setlogp!!(vi, state.loglikelihood)
+    vi = setloglikelihood!!(vi, state.loglikelihood)
 
     return Transition(model, vi), vi
 end
@@ -108,20 +108,12 @@ end
 # Mean of prior distribution
 Distributions.mean(p::ESSPrior) = p.μ
 
-# Evaluate log-likelihood of proposals
-const ESSLogLikelihood{M<:Model,S<:Sampler{<:ESS},V<:AbstractVarInfo} =
-    DynamicPPL.LogDensityFunction{M,V,<:DynamicPPL.SamplingContext{<:S},AD} where {AD}
-
-(ℓ::ESSLogLikelihood)(f::AbstractVector) = LogDensityProblems.logdensity(ℓ, f)
-
 function DynamicPPL.tilde_assume(
-    rng::Random.AbstractRNG, ::DefaultContext, ::Sampler{<:ESS}, right, vn, vi
+    rng::Random.AbstractRNG, ctx::DefaultContext, ::Sampler{<:ESS}, right, vn, vi
 )
-    return DynamicPPL.tilde_assume(
-        rng, LikelihoodContext(), SampleFromPrior(), right, vn, vi
-    )
+    return DynamicPPL.tilde_assume(rng, ctx, SampleFromPrior(), right, vn, vi)
 end
 
-function DynamicPPL.tilde_observe(ctx::DefaultContext, ::Sampler{<:ESS}, right, left, vi)
-    return DynamicPPL.tilde_observe(ctx, SampleFromPrior(), right, left, vi)
+function DynamicPPL.tilde_observe!!(ctx::DefaultContext, ::Sampler{<:ESS}, right, left, vn, vi)
+    return DynamicPPL.tilde_observe!!(ctx, SampleFromPrior(), right, left, vn, vi)
 end

src/mcmc/gibbs.jl

@@ -33,7 +33,7 @@ can_be_wrapped(ctx::DynamicPPL.PrefixContext) = can_be_wrapped(ctx.context)
 #
 # Purpose: avoid triggering resampling of variables we're conditioning on.
 # - Using standard `DynamicPPL.condition` results in conditioned variables being treated
-#   as observations in the truest sense, i.e. we hit `DynamicPPL.tilde_observe`.
+#   as observations in the truest sense, i.e. we hit `DynamicPPL.tilde_observe!!`.
 # - But `observe` is overloaded by some samplers, e.g. `CSMC`, which can lead to
 #   undesirable behavior, e.g. `CSMC` triggering a resampling for every conditioned variable
 #   rather than only for the "true" observations.
@@ -178,24 +178,26 @@ function DynamicPPL.tilde_assume(context::GibbsContext, right, vn, vi)
         DynamicPPL.tilde_assume(child_context, right, vn, vi)
     elseif has_conditioned_gibbs(context, vn)
         # Short-circuit the tilde assume if `vn` is present in `context`.
-        value, lp, _ = DynamicPPL.tilde_assume(
+        # TODO(mhauru) Fix accumulation here. In this branch anything that gets
+        # accumulated just gets discarded with `_`.
+        value, _ = DynamicPPL.tilde_assume(
             child_context, right, vn, get_global_varinfo(context)
         )
-        value, lp, vi
+        value, vi
     else
         # If the varname has not been conditioned on, nor is it a target variable, its
         # presumably a new variable that should be sampled from its prior. We need to add
         # this new variable to the global `varinfo` of the context, but not to the local one
         # being used by the current sampler.
-        value, lp, new_global_vi = DynamicPPL.tilde_assume(
+        value, new_global_vi = DynamicPPL.tilde_assume(
             child_context,
             DynamicPPL.SampleFromPrior(),
             right,
             vn,
             get_global_varinfo(context),
         )
         set_global_varinfo!(context, new_global_vi)
-        value, lp, vi
+        value, vi
     end
 end
 

src/mcmc/hmc.jl

@@ -516,10 +516,6 @@ function DynamicPPL.assume(
     return DynamicPPL.assume(dist, vn, vi)
 end
 
-function DynamicPPL.observe(::Sampler{<:Hamiltonian}, d::Distribution, value, vi)
-    return DynamicPPL.observe(d, value, vi)
-end
-
 ####
 #### Default HMC stepsize and mass matrix adaptor
 ####

src/mcmc/is.jl

@@ -55,7 +55,3 @@ function DynamicPPL.assume(rng, ::Sampler{<:IS}, dist::Distribution, vn::VarName
     end
     return r, 0, vi
 end
-
-function DynamicPPL.observe(::Sampler{<:IS}, dist::Distribution, value, vi)
-    return logpdf(dist, value), vi
-end

src/mcmc/mh.jl

@@ -392,7 +392,3 @@ function DynamicPPL.assume(
     retval = DynamicPPL.assume(rng, SampleFromPrior(), dist, vn, vi)
     return retval
 end
-
-function DynamicPPL.observe(spl::Sampler{<:MH}, d::Distribution, value, vi)
-    return DynamicPPL.observe(SampleFromPrior(), d, value, vi)
-end

src/mcmc/particle_mcmc.jl

@@ -450,10 +450,11 @@ function DynamicPPL.assume(
     return r, lp, vi
 end
 
-function DynamicPPL.observe(spl::Sampler{<:Union{PG,SMC}}, dist::Distribution, value, vi)
-    # NOTE: The `Libtask.produce` is now hit in `acclogp_observe!!`.
-    return logpdf(dist, value), trace_local_varinfo_maybe(vi)
-end
+# TODO(mhauru) Fix this.
+# function DynamicPPL.observe(spl::Sampler{<:Union{PG,SMC}}, dist::Distribution, value, vi)
+#     # NOTE: The `Libtask.produce` is now hit in `acclogp_observe!!`.
+#     return logpdf(dist, value), trace_local_varinfo_maybe(vi)
+# end
 
 function DynamicPPL.acclogp!!(
     context::SamplingContext{<:Sampler{<:Union{PG,SMC}}}, varinfo::AbstractVarInfo, logp
@@ -462,12 +463,13 @@ function DynamicPPL.acclogp!!(
     return DynamicPPL.acclogp!!(DynamicPPL.childcontext(context), varinfo_trace, logp)
 end
 
-function DynamicPPL.acclogp_observe!!(
-    context::SamplingContext{<:Sampler{<:Union{PG,SMC}}}, varinfo::AbstractVarInfo, logp
-)
-    Libtask.produce(logp)
-    return trace_local_varinfo_maybe(varinfo)
-end
+# TODO(mhauru) Fix this.
+# function DynamicPPL.acclogp_observe!!(
+#     context::SamplingContext{<:Sampler{<:Union{PG,SMC}}}, varinfo::AbstractVarInfo, logp
+# )
+#     Libtask.produce(logp)
+#     return trace_local_varinfo_maybe(varinfo)
+# end
 
 # Convenient constructor
 function AdvancedPS.Trace(