Fix Prior(), fix a couple more imports

Author: penelopeysm

src/mcmc/gibbs.jl

@@ -549,7 +549,8 @@ function setparams_varinfo!!(
     # update its logprob. To do this, we have to call evaluate!! with the sampler, rather
     # than just a context, because ESS is peculiar in how it uses LikelihoodContext for
     # some variables and DefaultContext for others.
-    return last(DynamicPPL.evaluate!!(model, params, SamplingContext(sampler)))
+    # TODO(penelopeysm): Is this still needed?
+    return last(DynamicPPL.evaluate!!(model, params, DynamicPPL.SamplingContext(sampler)))
 end
 
 function setparams_varinfo!!(
@@ -559,7 +560,7 @@ function setparams_varinfo!!(
     params::AbstractVarInfo,
 )
     logdensity = DynamicPPL.LogDensityFunction(
-        model, state.ldf.varinfo; adtype=sampler.alg.adtype
+        model, DynamicPPL.getlogjoint, state.ldf.varinfo; adtype=sampler.alg.adtype
     )
     new_inner_state = setparams_varinfo!!(
         AbstractMCMC.LogDensityModel(logdensity), sampler, state.state, params

src/mcmc/particle_mcmc.jl

@@ -18,7 +18,7 @@ function TracedModel(
     varinfo::AbstractVarInfo,
     rng::Random.AbstractRNG,
 )
-    context = SamplingContext(rng, sampler, DefaultContext())
+    context = DynamicPPL.SamplingContext(rng, sampler, DefaultContext())
     args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
     if kwargs !== nothing && !isempty(kwargs)
         error(
@@ -395,7 +395,7 @@ function AbstractMCMC.step(
 end
 
 function DynamicPPL.use_threadsafe_eval(
-    ::SamplingContext{<:Sampler{<:Union{PG,SMC}}}, ::AbstractVarInfo
+    ::DynamicPPL.SamplingContext{<:Sampler{<:Union{PG,SMC}}}, ::AbstractVarInfo
 )
     return false
 end
@@ -457,7 +457,9 @@ end
 # end
 
 function DynamicPPL.acclogp!!(
-    context::SamplingContext{<:Sampler{<:Union{PG,SMC}}}, varinfo::AbstractVarInfo, logp
+    context::DynamicPPL.SamplingContext{<:Sampler{<:Union{PG,SMC}}},
+    varinfo::AbstractVarInfo,
+    logp,
 )
     varinfo_trace = trace_local_varinfo_maybe(varinfo)
     return DynamicPPL.acclogp!!(DynamicPPL.childcontext(context), varinfo_trace, logp)

src/mcmc/prior.jl

@@ -12,14 +12,17 @@ function AbstractMCMC.step(
     state=nothing;
     kwargs...,
 )
+    # TODO(DPPL0.37/penelopeysm): replace with init!! instead
     vi = last(
         DynamicPPL.evaluate!!(
-            model,
-            VarInfo(),
-            SamplingContext(rng, DynamicPPL.SampleFromPrior(), DynamicPPL.PriorContext()),
+            model, VarInfo(), DynamicPPL.SamplingContext(rng, DynamicPPL.SampleFromPrior())
         ),
     )
-    return vi, nothing
+    # Need to manually construct the Transition here because we only
+    # want to use the prior probability.
+    xs = Turing.Inference.getparams(model, vi)
+    lp = DynamicPPL.getlogprior(vi)
+    return Transition(xs, lp, nothing)
 end
 
 DynamicPPL.default_chain_type(sampler::Prior) = MCMCChains.Chains