more fixes

Author: penelopeysm

src/model.jl

@@ -795,7 +795,7 @@ fixed(model::Model) = fixed(model.context)
 
 """
     (model::Model)()
-    (model::Model)(rng[, varinfo, sampler, context])
+    (model::Model)(rng[, varinfo])
 
 Sample from the `model` using the `sampler` with random number generator `rng`
 and the `context`, and store the sample and log joint probability in `varinfo`.
@@ -806,13 +806,8 @@ If no arguments are provided, uses the default random number generator and
 samples from the prior.
 """
 (model::Model)() = model(Random.default_rng())
-function (model::Model)(
-    rng::Random.AbstractRNG,
-    varinfo::AbstractVarInfo=VarInfo(),
-    sampler::AbstractSampler=SampleFromPrior(),
-)
-    spl_ctx = SamplingContext(rng, sampler, DefaultContext())
-    return first(evaluate!!(model, varinfo, spl_ctx))
+function (model::Model)(rng::Random.AbstractRNG, varinfo::AbstractVarInfo=VarInfo())
+    return first(sample!!(rng, model, varinfo))
 end
 
 """
@@ -1016,7 +1011,7 @@ Base.nameof(model::Model{<:Function}) = nameof(model.f)
 Generate a sample of type `T` from the prior distribution of the `model`.
 """
 function Base.rand(rng::Random.AbstractRNG, ::Type{T}, model::Model) where {T}
-    x = last(sample!!(model, SimpleVarInfo{Float64}(OrderedDict())))
+    x = last(sample!!(rng, model, SimpleVarInfo{Float64}(OrderedDict())))
     return values_as(x, T)
 end
 
@@ -1087,7 +1082,7 @@ function logprior(model::Model, varinfo::AbstractVarInfo)
         LogPriorAccumulator()
     end
     varinfo = setaccs!!(deepcopy(varinfo), (logprioracc,))
-    return getlogprior(last(evaluate!!(model, varinfo, DefaultContext())))
+    return getlogprior(last(evaluate!!(model, varinfo)))
 end
 
 """
@@ -1141,7 +1136,7 @@ function Distributions.loglikelihood(model::Model, varinfo::AbstractVarInfo)
         LogLikelihoodAccumulator()
     end
     varinfo = setaccs!!(deepcopy(varinfo), (loglikelihoodacc,))
-    return getloglikelihood(last(evaluate!!(model, varinfo, DefaultContext())))
+    return getloglikelihood(last(evaluate!!(model, varinfo)))
 end
 
 """
@@ -1195,7 +1190,7 @@ function predict(
     return map(chain) do params_varinfo
         vi = deepcopy(varinfo)
         DynamicPPL.setval_and_resample!(vi, values_as(params_varinfo, NamedTuple))
-        model(rng, vi, SampleFromPrior())
+        model(rng, vi)
         return vi
     end
 end

src/threadsafe.jl

@@ -116,14 +116,17 @@ end
 # consistency between `vi.accs_by_thread` field and `getacc(vi.varinfo)`, which accumulates
 # to define `getacc(vi)`.
 function link!!(t::DynamicTransformation, vi::ThreadSafeVarInfo, model::Model)
-    return settrans!!(last(evaluate!!(model, vi, DynamicTransformationContext{false}())), t)
+    model = contextualize(
+        model, setleafcontext(model.context, DynamicTransformationContext{false}())
+    )
+    return settrans!!(last(evaluate!!(model, vi)), t)
 end
 
 function invlink!!(::DynamicTransformation, vi::ThreadSafeVarInfo, model::Model)
-    return settrans!!(
-        last(evaluate!!(model, vi, DynamicTransformationContext{true}())),
-        NoTransformation(),
+    model = contextualize(
+        model, setleafcontext(model.context, DynamicTransformationContext{true}())
     )
+    return settrans!!(last(evaluate!!(model, vi)), NoTransformation())
 end
 
 function link(t::DynamicTransformation, vi::ThreadSafeVarInfo, model::Model)

src/transforming.jl

@@ -51,16 +51,17 @@ function _transform!!(
     vi::AbstractVarInfo,
     model::Model,
 )
-    # To transform using DynamicTransformationContext, we evaluate the model, but we do not
-    # need to use any accumulators other than LogPriorAccumulator (which is affected by the Jacobian of
-    # the transformation).
+    # To transform using DynamicTransformationContext, we evaluate the model using that as the leaf context:
+    model = contextualize(model, setleafcontext(model.context, ctx))
+    # but we do not need to use any accumulators other than LogPriorAccumulator
+    # (which is affected by the Jacobian of the transformation).
     accs = getaccs(vi)
     has_logprior = haskey(accs, Val(:LogPrior))
     if has_logprior
         old_logprior = getacc(accs, Val(:LogPrior))
         vi = setaccs!!(vi, (old_logprior,))
     end
-    vi = settrans!!(last(evaluate!!(model, vi, ctx)), t)
+    vi = settrans!!(last(evaluate!!(model, vi)), t)
     # Restore the accumulators.
     if has_logprior
         new_logprior = getacc(vi, Val(:LogPrior))

test/model.jl

@@ -332,8 +332,7 @@ const GDEMO_DEFAULT = DynamicPPL.TestUtils.demo_assume_observe_literal()
             @test logjoint(model, x) !=
                 DynamicPPL.TestUtils.logjoint_true_with_logabsdet_jacobian(model, x...)
             # Ensure `varnames` is implemented.
-            sampling_model = contextualize(model, SamplingContext(model.context))
-            vi = last(DynamicPPL.evaluate!!(sampling_model, SimpleVarInfo(OrderedDict())))
+            vi = last(DynamicPPL.sample!!(sampling_model, SimpleVarInfo(OrderedDict())))
             @test all(collect(keys(vi)) .== DynamicPPL.TestUtils.varnames(model))
             # Ensure `posterior_mean` is implemented.
             @test DynamicPPL.TestUtils.posterior_mean(model) isa typeof(x)