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

@@ -0,0 +1,73 @@
+struct TracedModel{S<:AbstractSampler,V<:AbstractVarInfo,M<:Model,E<:Tuple} <:
+       AdvancedPS.AbstractGenericModel
+    model::M
+    sampler::S
+    varinfo::V
+    evaluator::E
+end
+
+function TracedModel(
+    model::Model,
+    sampler::AbstractSampler,
+    varinfo::AbstractVarInfo,
+    rng::Random.AbstractRNG,
+)
+    context = SamplingContext(rng, sampler, DefaultContext())
+    args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
+    if kwargs !== nothing && !isempty(kwargs)
+        error(
+            "Sampling with `$(sampler.alg)` does not support models with keyword arguments. See issue #2007 for more details.",
+        )
+    end
+    return TracedModel{AbstractSampler,AbstractVarInfo,Model,Tuple}(
+        model, sampler, varinfo, (model.f, args...)
+    )
+end
+
+function AdvancedPS.advance!(
+    trace::AdvancedPS.Trace{<:AdvancedPS.LibtaskModel{<:TracedModel}}, isref::Bool=false
+)
+    # Make sure we load/reset the rng in the new replaying mechanism
+    # 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
+        return nothing
+    else
+        return score + DynamicPPL.getlogp(trace.model.f.varinfo)
+    end
+end
+
+function AdvancedPS.delete_retained!(trace::TracedModel)
+    DynamicPPL.set_retained_vns_del!(trace.varinfo)
+    return trace
+end
+
+function AdvancedPS.reset_model(trace::TracedModel)
+    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
+
+function AdvancedPS.update_rng!(
+    trace::AdvancedPS.Trace{<:AdvancedPS.LibtaskModel{<:TracedModel}}
+)
+    # Extract the `args`.
+    args = trace.model.ctask.args
+    # From `args`, extract the `SamplingContext`, which contains the RNG.
+    sampling_context = args[3]
+    rng = sampling_context.rng
+    trace.rng = rng
+    return trace
+end
+
+function Libtask.TapedTask(model::TracedModel, ::Random.AbstractRNG, args...; kwargs...) # RNG ?
+    return Libtask.TapedTask(model.evaluator[1], model.evaluator[2:end]...; kwargs...)
+end

src/mcmc/Inference.jl

@@ -18,6 +18,7 @@ using DynamicPPL:
     push!!,
     setlogp!!,
     getlogp,
+    getlogjoint,
     VarName,
     getsym,
     getdist,
@@ -136,7 +137,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
 
@@ -149,10 +150,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),)
 
 ##########################
 # Chain making utilities #
@@ -215,7 +216,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
 
@@ -434,7 +435,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

@@ -214,7 +214,7 @@ function DynamicPPL.initialstep(
     theta = vi[:]
 
     # 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.ϵ)
@@ -242,10 +242,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)
@@ -290,7 +292,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

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