More DPPL 0.37 compat work, WIP

Author: mhauru

ext/TuringOptimExt.jl

@@ -34,8 +34,7 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    vi = DynamicPPL.setaccs!!(VarInfo(model), (DynamicPPL.LogLikelihoodAccumulator(),))
-    f = Optimisation.OptimLogDensity(model, vi)
+    f = Optimisation.OptimLogDensity(model, DynamicPPL.getloglikelihood)
     init_vals = DynamicPPL.getparams(f.ldf)
     optimizer = Optim.LBFGS()
     return _mle_optimize(model, init_vals, optimizer, options; kwargs...)
@@ -57,8 +56,7 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    vi = DynamicPPL.setaccs!!(VarInfo(model), (DynamicPPL.LogLikelihoodAccumulator(),))
-    f = Optimisation.OptimLogDensity(model, vi)
+    f = Optimisation.OptimLogDensity(model, DynamicPPL.getloglikelihood)
     init_vals = DynamicPPL.getparams(f.ldf)
     return _mle_optimize(model, init_vals, optimizer, options; kwargs...)
 end
@@ -74,8 +72,7 @@ function Optim.optimize(
 end
 
 function _mle_optimize(model::DynamicPPL.Model, args...; kwargs...)
-    vi = DynamicPPL.setaccs!!(VarInfo(model), (DynamicPPL.LogLikelihoodAccumulator(),))
-    f = Optimisation.OptimLogDensity(model, vi)
+    f = Optimisation.OptimLogDensity(model, DynamicPPL.getloglikelihood)
     return _optimize(f, args...; kwargs...)
 end
 
@@ -105,8 +102,7 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    vi = DynamicPPL.setaccs!!(VarInfo(model), (LogPriorWithoutJacobianAccumulator(), DynamicPPL.LogLikelihoodAccumulator(),))
-    f = Optimisation.OptimLogDensity(model, vi)
+    f = Optimisation.OptimLogDensity(model, Optimisation.getlogjoint_without_jacobian)
     init_vals = DynamicPPL.getparams(f.ldf)
     optimizer = Optim.LBFGS()
     return _map_optimize(model, init_vals, optimizer, options; kwargs...)
@@ -128,8 +124,7 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    vi = DynamicPPL.setaccs!!(VarInfo(model), (LogPriorWithoutJacobianAccumulator(), DynamicPPL.LogLikelihoodAccumulator(),))
-    f = Optimisation.OptimLogDensity(model, vi)
+    f = Optimisation.OptimLogDensity(model, Optimisation.getlogjoint_without_jacobian)
     init_vals = DynamicPPL.getparams(f.ldf)
     return _map_optimize(model, init_vals, optimizer, options; kwargs...)
 end
@@ -145,8 +140,7 @@ function Optim.optimize(
 end
 
 function _map_optimize(model::DynamicPPL.Model, args...; kwargs...)
-    vi = DynamicPPL.setaccs!!(VarInfo(model), (LogPriorWithoutJacobianAccumulator(), DynamicPPL.LogLikelihoodAccumulator(),))
-    f = Optimisation.OptimLogDensity(model, vi)
+    f = Optimisation.OptimLogDensity(model, Optimisation.getlogjoint_without_jacobian)
     return _optimize(f, args...; kwargs...)
 end
 
@@ -169,7 +163,9 @@ 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; adtype=f.ldf.adtype)
+    f = Optimisation.OptimLogDensity(
+        f.ldf.model, f.ldf.getlogdensity, vi; adtype=f.ldf.adtype
+    )
     init_vals = DynamicPPL.getparams(f.ldf)
 
     # Optimize!
@@ -186,7 +182,9 @@ 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; adtype=f.ldf.adtype)
+    logdensity_optimum = Optimisation.OptimLogDensity(
+        f.ldf.model, f.ldf.getlogdensity, 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/essential/container.jl

@@ -28,7 +28,8 @@ function AdvancedPS.advance!(
     trace::AdvancedPS.Trace{<:AdvancedPS.LibtaskModel{<:TracedModel}}, isref::Bool=false
 )
     # Make sure we load/reset the rng in the new replaying mechanism
-    DynamicPPL.increment_num_produce!(trace.model.f.varinfo)
+    # TODO(mhauru) Stop ignoring the return value.
+    DynamicPPL.increment_num_produce!!(trace.model.f.varinfo)
     isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.rng)
     score = consume(trace.model.ctask)
     if score === nothing
@@ -44,11 +45,13 @@ function AdvancedPS.delete_retained!(trace::TracedModel)
 end
 
 function AdvancedPS.reset_model(trace::TracedModel)
-    DynamicPPL.reset_num_produce!(trace.varinfo)
+    new_vi = DynamicPPL.reset_num_produce!!(trace.varinfo)
+    trace = TracedModel(trace.model, trace.sampler, new_vi, trace.evaluator)
     return trace
 end
 
 function AdvancedPS.reset_logprob!(trace::TracedModel)
+    # TODO(mhauru) Stop ignoring the return value.
     DynamicPPL.resetlogp!!(trace.model.varinfo)
     return trace
 end

src/mcmc/Inference.jl

@@ -14,6 +14,7 @@ using DynamicPPL:
     push!!,
     setlogp!!,
     getlogp,
+    getlogjoint,
     VarName,
     getsym,
     getdist,
@@ -182,7 +183,7 @@ function AbstractMCMC.step(
     vi = VarInfo()
     vi = DynamicPPL.setaccs!!(vi, (DynamicPPL.LogPriorAccumulator(),))
     vi = last(
-        DynamicPPL.evaluate!!(model, vi, SamplingContext(rng, DynamicPPL.SampleFromPrior())),
+        DynamicPPL.evaluate!!(model, vi, SamplingContext(rng, DynamicPPL.SampleFromPrior()))
     )
     return vi, nothing
 end
@@ -223,7 +224,7 @@ end
 Transition(θ, lp) = Transition(θ, lp, nothing)
 function Transition(model::DynamicPPL.Model, vi::AbstractVarInfo, t)
     θ = getparams(model, vi)
-    lp = getlogp(vi)
+    lp = getlogjoint(vi)
     return Transition(θ, lp, getstats(t))
 end
 
@@ -236,10 +237,10 @@ function metadata(t::Transition)
     end
 end
 
-DynamicPPL.getlogp(t::Transition) = t.lp
+DynamicPPL.getlogjoint(t::Transition) = t.lp
 
 # Metadata of VarInfo object
-metadata(vi::AbstractVarInfo) = (lp=getlogp(vi),)
+metadata(vi::AbstractVarInfo) = (lp=getlogjoint(vi),)
 
 # TODO: Implement additional checks for certain samplers, e.g.
 # HMC not supporting discrete parameters.
@@ -376,7 +377,7 @@ function _params_to_array(model::DynamicPPL.Model, ts::Vector)
 end
 
 function get_transition_extras(ts::AbstractVector{<:VarInfo})
-    valmat = reshape([getlogp(t) for t in ts], :, 1)
+    valmat = reshape([getlogjoint(t) for t in ts], :, 1)
     return [:lp], valmat
 end
 
@@ -589,7 +590,7 @@ julia> chain = Chains(randn(2, 1, 1), ["m"]); # 2 samples of `m`
 
 julia> transitions = Turing.Inference.transitions_from_chain(m, chain);
 
-julia> [Turing.Inference.getlogp(t) for t in transitions] # extract the logjoints
+julia> [Turing.Inference.getlogjoint(t) for t in transitions] # extract the logjoints
 2-element Array{Float64,1}:
  -3.6294991938628374
  -2.5697948166987845

src/mcmc/ess.jl

@@ -114,6 +114,8 @@ function DynamicPPL.tilde_assume(
     return DynamicPPL.tilde_assume(rng, ctx, SampleFromPrior(), right, vn, vi)
 end
 
-function DynamicPPL.tilde_observe!!(ctx::DefaultContext, ::Sampler{<:ESS}, right, left, vn, 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/hmc.jl

@@ -199,7 +199,7 @@ function DynamicPPL.initialstep(
     end
 
     # Cache current log density.
-    log_density_old = getlogp(vi)
+    log_density_old = getloglikelihood(vi)
 
     # Find good eps if not provided one
     if iszero(spl.alg.ϵ)
@@ -227,10 +227,12 @@ function DynamicPPL.initialstep(
     # Update `vi` based on acceptance
     if t.stat.is_accept
         vi = DynamicPPL.unflatten(vi, t.z.θ)
-        vi = setlogp!!(vi, t.stat.log_density)
+        # TODO(mhauru) Is setloglikelihood! the right thing here?
+        vi = setloglikelihood!!(vi, t.stat.log_density)
     else
         vi = DynamicPPL.unflatten(vi, theta)
-        vi = setlogp!!(vi, log_density_old)
+        # TODO(mhauru) Is setloglikelihood! the right thing here?
+        vi = setloglikelihood!!(vi, log_density_old)
     end
 
     transition = Transition(model, vi, t)
@@ -275,7 +277,8 @@ function AbstractMCMC.step(
     vi = state.vi
     if t.stat.is_accept
         vi = DynamicPPL.unflatten(vi, t.z.θ)
-        vi = setlogp!!(vi, t.stat.log_density)
+        # TODO(mhauru) Is setloglikelihood! the right thing here?
+        vi = setloglikelihood!!(vi, t.stat.log_density)
     end
 
     # Compute next transition and state.

src/mcmc/particle_mcmc.jl

@@ -118,10 +118,10 @@ function DynamicPPL.initialstep(
     kwargs...,
 )
     # Reset the VarInfo.
-    reset_num_produce!(vi)
+    vi = reset_num_produce!!(vi)
     set_retained_vns_del!(vi)
-    resetlogp!!(vi)
-    empty!!(vi)
+    vi = resetlogp!!(vi)
+    vi = empty!!(vi)
 
     # Create a new set of particles.
     particles = AdvancedPS.ParticleContainer(
@@ -252,9 +252,9 @@ function DynamicPPL.initialstep(
     kwargs...,
 )
     # Reset the VarInfo before new sweep
-    reset_num_produce!(vi)
+    vi = reset_num_produce!(vi)
     set_retained_vns_del!(vi)
-    resetlogp!!(vi)
+    vi = resetlogp!!(vi)
 
     # Create a new set of particles
     num_particles = spl.alg.nparticles
@@ -284,8 +284,8 @@ function AbstractMCMC.step(
 )
     # Reset the VarInfo before new sweep.
     vi = state.vi
-    reset_num_produce!(vi)
-    resetlogp!!(vi)
+    vi = reset_num_produce!(vi)
+    vi = resetlogp!!(vi)
 
     # Create reference particle for which the samples will be retained.
     reference = AdvancedPS.forkr(AdvancedPS.Trace(model, spl, vi, state.rng))
@@ -408,7 +408,7 @@ function AdvancedPS.Trace(
     rng::AdvancedPS.TracedRNG,
 )
     newvarinfo = deepcopy(varinfo)
-    DynamicPPL.reset_num_produce!(newvarinfo)
+    newvarinfo = DynamicPPL.reset_num_produce!!(newvarinfo)
 
     tmodel = Turing.Essential.TracedModel(model, sampler, newvarinfo, rng)
     newtrace = AdvancedPS.Trace(tmodel, rng)