more fixes

Author: penelopeysm

ext/DynamicPPLJETExt.jl

@@ -31,7 +31,7 @@ function DynamicPPL.Experimental._determine_varinfo_jet(
 
     # Let's make sure that both evaluation and sampling doesn't result in type errors.
     issuccess, result = DynamicPPL.Experimental.is_suitable_varinfo(
-        model, varinfo; only_ddpl
+        sampling_model, varinfo; only_ddpl
     )
 
     if !issuccess
@@ -46,7 +46,7 @@ function DynamicPPL.Experimental._determine_varinfo_jet(
     else
         # Warn the user that we can't use the type stable one.
         @warn "Model seems incompatible with typed varinfo. Falling back to untyped varinfo."
-        DynamicPPL.untyped_varinfo(model)
+        DynamicPPL.untyped_varinfo(sampling_model)
     end
 end
 

ext/DynamicPPLMCMCChainsExt.jl

@@ -115,7 +115,7 @@ function DynamicPPL.predict(
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
     predictive_samples = map(iters) do (sample_idx, chain_idx)
         DynamicPPL.setval_and_resample!(varinfo, parameter_only_chain, sample_idx, chain_idx)
-        model(rng, varinfo, DynamicPPL.SampleFromPrior())
+        varinfo = last(DynamicPPL.sample!!(rng, model, varinfo))
 
         vals = DynamicPPL.values_as_in_model(model, false, varinfo)
         varname_vals = mapreduce(

src/logdensityfunction.jl

@@ -18,7 +18,7 @@ is_supported(::ADTypes.AutoReverseDiff) = true
 """
     LogDensityFunction(
         model::Model,
-        varinfo::AbstractVarInfo=VarInfo(model),
+        varinfo::AbstractVarInfo=VarInfo(model);
         adtype::Union{ADTypes.AbstractADType,Nothing}=nothing
     )
 
@@ -106,7 +106,7 @@ struct LogDensityFunction{
 
     function LogDensityFunction(
         model::Model,
-        varinfo::AbstractVarInfo=VarInfo(model),
+        varinfo::AbstractVarInfo=VarInfo(model);
         adtype::Union{ADTypes.AbstractADType,Nothing}=nothing,
     )
         if adtype === nothing

src/model.jl

@@ -794,8 +794,7 @@ julia> # Now `a.x` will be sampled.
 fixed(model::Model) = fixed(model.context)
 
 """
-    (model::Model)()
-    (model::Model)(rng[, varinfo])
+    (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`.
@@ -805,10 +804,12 @@ Returns the model's return value.
 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())
     return first(sample!!(rng, model, varinfo))
 end
+function (model::Model)(varinfo::AbstractVarInfo=VarInfo())
+    return model(Random.default_rng(), varinfo)
+end
 
 """
     use_threadsafe_eval(context::AbstractContext, varinfo::AbstractVarInfo)
@@ -821,21 +822,29 @@ function use_threadsafe_eval(context::AbstractContext, varinfo::AbstractVarInfo)
 end
 
 """
-    sample!!([rng::Random.AbstractRNG, ]model::Model, varinfo)
+    sample!!([rng::Random.AbstractRNG, ]model::Model, varinfo[, sampler])
 
 Evaluate the `model` with the given `varinfo`, but perform sampling during the
-evaluation by wrapping the model's context in a `SamplingContext`.
+evaluation using the given `sampler` by wrapping the model's context in a
+`SamplingContext`.
+
+If `sampler` is not provided, defaults to [`SampleFromPrior`](@ref).
 
 Returns a tuple of the model's return value, plus the updated `varinfo` object.
 """
-function sample!!(rng::Random.AbstractRNG, model::Model, varinfo::AbstractVarInfo)
-    sampling_model = contextualize(
-        model, SamplingContext(rng, SampleFromPrior(), model.context)
-    )
+function sample!!(
+    rng::Random.AbstractRNG,
+    model::Model,
+    varinfo::AbstractVarInfo,
+    sampler::AbstractSampler=SampleFromPrior(),
+)
+    sampling_model = contextualize(model, SamplingContext(rng, sampler, model.context))
     return evaluate!!(sampling_model, varinfo)
 end
-function sample!!(model::Model, varinfo::AbstractVarInfo)
-    return sample!!(Random.default_rng(), model, varinfo)
+function sample!!(
+    model::Model, varinfo::AbstractVarInfo, sampler::AbstractSampler=SampleFromPrior()
+)
+    return sample!!(Random.default_rng(), model, varinfo, sampler)
 end
 
 """
@@ -1028,7 +1037,7 @@ Return the log joint probability of variables `varinfo` for the probabilistic `m
 See [`logprior`](@ref) and [`loglikelihood`](@ref).
 """
 function logjoint(model::Model, varinfo::AbstractVarInfo)
-    return getlogjoint(last(evaluate!!(model, varinfo, DefaultContext())))
+    return getlogjoint(last(evaluate!!(model, varinfo)))
 end
 
 """

src/test_utils/ad.jl

@@ -60,8 +60,6 @@ struct ADResult{Tparams<:AbstractFloat,Tresult<:AbstractFloat}
     model::Model
     "The VarInfo that was used"
     varinfo::AbstractVarInfo
-    "The evaluation context that was used"
-    context::AbstractContext
     "The values at which the model was evaluated"
     params::Vector{Tparams}
     "The AD backend that was tested"
@@ -92,7 +90,6 @@ end
         grad_atol=1e-6,
         varinfo::AbstractVarInfo=link(VarInfo(model), model),
         params::Union{Nothing,Vector{<:AbstractFloat}}=nothing,
-        context::AbstractContext=DefaultContext(),
         reference_adtype::ADTypes.AbstractADType=REFERENCE_ADTYPE,
         expected_value_and_grad::Union{Nothing,Tuple{AbstractFloat,Vector{<:AbstractFloat}}}=nothing,
         verbose=true,
@@ -146,13 +143,7 @@ Everything else is optional, and can be categorised into several groups:
    prep_params)`. You could then evaluate the gradient at a different set of
    parameters using the `params` keyword argument.
 
-3. _How to specify the evaluation context._
-
-   A `DynamicPPL.AbstractContext` can be passed as the `context` keyword
-   argument to control the evaluation context. This defaults to
-   `DefaultContext()`.
-
-4. _How to specify the results to compare against._ (Only if `test=true`.)
+3. _How to specify the results to compare against._ (Only if `test=true`.)
 
    Once logp and its gradient has been calculated with the specified `adtype`,
    it must be tested for correctness.
@@ -167,12 +158,12 @@ Everything else is optional, and can be categorised into several groups:
    The default reference backend is ForwardDiff. If none of these parameters are
    specified, ForwardDiff will be used to calculate the ground truth.
 
-5. _How to specify the tolerances._ (Only if `test=true`.)
+4. _How to specify the tolerances._ (Only if `test=true`.)
 
    The tolerances for the value and gradient can be set using `value_atol` and
    `grad_atol`. These default to 1e-6.
 
-6. _Whether to output extra logging information._
+5. _Whether to output extra logging information._
 
    By default, this function prints messages when it runs. To silence it, set
    `verbose=false`.
@@ -195,7 +186,6 @@ function run_ad(
     grad_atol::AbstractFloat=1e-6,
     varinfo::AbstractVarInfo=link(VarInfo(model), model),
     params::Union{Nothing,Vector{<:AbstractFloat}}=nothing,
-    context::AbstractContext=DefaultContext(),
     reference_adtype::AbstractADType=REFERENCE_ADTYPE,
     expected_value_and_grad::Union{Nothing,Tuple{AbstractFloat,Vector{<:AbstractFloat}}}=nothing,
     verbose=true,
@@ -207,7 +197,7 @@ function run_ad(
 
     verbose && @info "Running AD on $(model.f) with $(adtype)\n"
     verbose && println("       params : $(params)")
-    ldf = LogDensityFunction(model, varinfo, context; adtype=adtype)
+    ldf = LogDensityFunction(model, varinfo; adtype=adtype)
 
     value, grad = logdensity_and_gradient(ldf, params)
     grad = collect(grad)
@@ -216,7 +206,7 @@ function run_ad(
     if test
         # Calculate ground truth to compare against
         value_true, grad_true = if expected_value_and_grad === nothing
-            ldf_reference = LogDensityFunction(model, varinfo, context; adtype=reference_adtype)
+            ldf_reference = LogDensityFunction(model, varinfo; adtype=reference_adtype)
             logdensity_and_gradient(ldf_reference, params)
         else
             expected_value_and_grad
@@ -245,7 +235,6 @@ function run_ad(
     return ADResult(
         model,
         varinfo,
-        context,
         params,
         adtype,
         value_atol,

test/compiler.jl

@@ -185,10 +185,7 @@ module Issue537 end
         @model function testmodel_missing3(x)
             x[1] ~ Bernoulli(0.5)
             global varinfo_ = __varinfo__
-            global sampler_ = __context__.sampler
             global model_ = __model__
-            global context_ = __context__
-            global rng_ = __context__.rng
             global lp = getlogjoint(__varinfo__)
             return x
         end
@@ -197,17 +194,14 @@ module Issue537 end
         @test getlogjoint(varinfo) == lp
         @test varinfo_ isa AbstractVarInfo
         @test model_ === model
-        @test context_ isa SamplingContext
-        @test rng_ isa Random.AbstractRNG
+        @test model_.context isa SamplingContext
+        @test model_.context.rng isa Random.AbstractRNG
 
         # disable warnings
         @model function testmodel_missing4(x)
             x[1] ~ Bernoulli(0.5)
             global varinfo_ = __varinfo__
-            global sampler_ = __context__.sampler
             global model_ = __model__
-            global context_ = __context__
-            global rng_ = __context__.rng
             global lp = getlogjoint(__varinfo__)
             return x
         end false

test/ext/DynamicPPLJETExt.jl

@@ -62,6 +62,7 @@
 
     @testset "demo models" begin
         @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
+            sampling_model = contextualize(model, SamplingContext(model.context))
             # Use debug logging below.
             varinfo = DynamicPPL.Experimental.determine_suitable_varinfo(model)
             # Check that the inferred varinfo is indeed suitable for evaluation and sampling
@@ -71,7 +72,7 @@
             JET.test_call(f_eval, argtypes_eval)
 
             f_sample, argtypes_sample = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
-                model, varinfo, DynamicPPL.SamplingContext()
+                sampling_model, varinfo
             )
             JET.test_call(f_sample, argtypes_sample)
             # For our demo models, they should all result in typed.
@@ -85,7 +86,7 @@
                 )
                 JET.test_call(f_eval, argtypes_eval)
                 f_sample, argtypes_sample = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
-                    model, typed_vi, DynamicPPL.SamplingContext()
+                    sampling_model, typed_vi
                 )
                 JET.test_call(f_sample, argtypes_sample)
             end

test/varinfo.jl

@@ -491,7 +491,7 @@ end
         # Check that instantiating the model does not perform linking
         vi = VarInfo()
         meta = vi.metadata
-        model(vi, SampleFromUniform())
+        model(vi)
         @test all(x -> !istrans(vi, x), meta.vns)
 
         # Check that linking and invlinking set the `trans` flag accordingly