fix

Author: penelopeysm

src/pointwise_logdensities.jl

@@ -3,9 +3,10 @@
 
 An accumulator that stores the log-probabilities of each variable in a model.
 
-Internally this context stores the log-probabilities in a dictionary, where the keys are
-the variable names and the values are vectors of log-probabilities. Each element in a vector
-corresponds to one execution of the model.
+Internally this accumulator stores the log-probabilities in a dictionary, where
+the keys are the variable names and the values are vectors of
+log-probabilities. Each element in a vector corresponds to one execution of the
+model.
 
 `whichlogprob` is a symbol that can be `:both`, `:prior`, or `:likelihood`, and specifies
 which log-probabilities to store in the accumulator. `KeyType` is the type by which variable
@@ -258,10 +259,8 @@ function pointwise_loglikelihoods(model::Model, chain, keytype::Type{T}=String)
     return pointwise_logdensities(model, chain, T, Val(:likelihood))
 end
 
-function pointwise_loglikelihoods(
-    model::Model, varinfo::AbstractVarInfo, context::AbstractContext=DefaultContext()
-)
-    return pointwise_logdensities(model, varinfo, context, Val(:likelihood))
+function pointwise_loglikelihoods(model::Model, varinfo::AbstractVarInfo)
+    return pointwise_logdensities(model, varinfo, Val(:likelihood))
 end
 
 """

src/sampler.jl

@@ -63,30 +63,21 @@ function AbstractMCMC.step(
 end
 
 """
-    default_varinfo(rng, model, sampler[, context])
+    default_varinfo(rng, model, sampler)
 
 Return a default varinfo object for the given `model` and `sampler`.
 
 # Arguments
 - `rng::Random.AbstractRNG`: Random number generator.
 - `model::Model`: Model for which we want to create a varinfo object.
 - `sampler::AbstractSampler`: Sampler which will make use of the varinfo object.
-- `context::AbstractContext`: Context in which the model is evaluated.
 
 # Returns
 - `AbstractVarInfo`: Default varinfo object for the given `model` and `sampler`.
 """
 function default_varinfo(rng::Random.AbstractRNG, model::Model, sampler::AbstractSampler)
-    return default_varinfo(rng, model, sampler, DefaultContext())
-end
-function default_varinfo(
-    rng::Random.AbstractRNG,
-    model::Model,
-    sampler::AbstractSampler,
-    context::AbstractContext,
-)
     init_sampler = initialsampler(sampler)
-    return typed_varinfo(rng, model, init_sampler, context)
+    return typed_varinfo(rng, model, init_sampler)
 end
 
 function AbstractMCMC.sample(

test/model.jl

@@ -591,7 +591,7 @@ const GDEMO_DEFAULT = DynamicPPL.TestUtils.demo_assume_observe_literal()
             xs_train = 1:0.1:10
             ys_train = ground_truth_β .* xs_train + rand(Normal(0, 0.1), length(xs_train))
             m_lin_reg = linear_reg(xs_train, ys_train)
-            chain = [evaluate!!(m_lin_reg)[2] for _ in 1:10000]
+            chain = [last(evaluate!!(m_lin_reg, VarInfo())) for _ in 1:10000]
 
             # chain is generated from the prior
             @test mean([chain[i][@varname(β)] for i in eachindex(chain)]) ≈ 1.0 atol = 0.1