InitContext, part 4 - Use init!! to replace evaluate_and_sample!!, predict, returned, and initialize_values (#984)

* Replace `evaluate_and_sample!!` -> `init!!`

* Use `ParamsInit` for `predict`; remove `setval_and_resample!` and friends

* Use `init!!` for initialisation

* Paper over the `Sampling->Init` context stack (pending removal of SamplingContext)

* Remove SamplingContext from JETExt to avoid triggering `Sampling->Init` pathway

* Remove `predict` on vector of VarInfo

* Fix some tests

* Remove duplicated test

* Simplify context testing

* Rename FooInit -> InitFromFoo

* Fix JETExt

* Fix JETExt properly

* Fix tests

* Improve comments

* Remove duplicated tests

* Docstring improvements

Co-authored-by: Markus Hauru <[email protected]>

* Concretise `chain_sample_to_varname_dict` using chain value type

* Clarify testset name

* Re-add comment that shouldn't have vanished

* Fix stale Requires dep

* Fix default_varinfo/initialisation for odd models

* Add comment to src/sampler.jl

Co-authored-by: Markus Hauru <[email protected]>

---------

Co-authored-by: Markus Hauru <[email protected]>

Author: penelopeysm

Project.toml

@@ -23,7 +23,6 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
@@ -71,7 +70,6 @@ Mooncake = "0.4.147"
 OrderedCollections = "1"
 Printf = "1.10"
 Random = "1.6"
-Requires = "1"
 Statistics = "1"
 Test = "1.6"
 julia = "1.10.8"

docs/src/api.md

@@ -447,11 +447,6 @@ AbstractPPL.evaluate!!
 
 This method mutates the `varinfo` used for execution.
 By default, it does not perform any actual sampling: it only evaluates the model using the values of the variables that are already in the `varinfo`.
-To perform sampling, you can either wrap `model.context` in a `SamplingContext`, or use this convenience method:
-
-```@docs
-DynamicPPL.evaluate_and_sample!!
-```
 
 The behaviour of a model execution can be changed with evaluation contexts, which are a field of the model.
 Contexts are subtypes of `AbstractPPL.AbstractContext`.
@@ -466,7 +461,12 @@ InitContext
 
 ### VarInfo initialisation
 
-`InitContext` is used to initialise, or overwrite, values in a VarInfo.
+The function `init!!` is used to initialise, or overwrite, values in a VarInfo.
+It is really a thin wrapper around using `evaluate!!` with an `InitContext`.
+
+```@docs
+DynamicPPL.init!!
+```
 
 To accomplish this, an initialisation _strategy_ is required, which defines how new values are to be obtained.
 There are three concrete strategies provided in DynamicPPL:
@@ -505,7 +505,7 @@ The default implementation of [`Sampler`](@ref) uses the following unexported fu
 ```@docs
 DynamicPPL.initialstep
 DynamicPPL.loadstate
-DynamicPPL.initialsampler
+DynamicPPL.init_strategy
 ```
 
 Finally, to specify which varinfo type a [`Sampler`](@ref) should use for a given [`Model`](@ref), this is specified by [`DynamicPPL.default_varinfo`](@ref) and can thus be overloaded for each  `model`-`sampler` combination. This can be useful in cases where one has explicit knowledge that one type of varinfo will be more performant for the given `model` and `sampler`.

ext/DynamicPPLJETExt.jl

@@ -6,7 +6,6 @@ using JET: JET
 function DynamicPPL.Experimental.is_suitable_varinfo(
     model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo; only_ddpl::Bool=true
 )
-    # Let's make sure that both evaluation and sampling doesn't result in type errors.
     f, argtypes = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(model, varinfo)
     # If specified, we only check errors originating somewhere in the DynamicPPL.jl.
     # This way we don't just fall back to untyped if the user's code is the issue.
@@ -21,32 +20,40 @@ end
 function DynamicPPL.Experimental._determine_varinfo_jet(
     model::DynamicPPL.Model; only_ddpl::Bool=true
 )
-    # Use SamplingContext to test type stability.
-    sampling_model = DynamicPPL.contextualize(
-        model, DynamicPPL.SamplingContext(model.context)
-    )
-
-    # First we try with the typed varinfo.
-    varinfo = DynamicPPL.typed_varinfo(sampling_model)
+    # Generate a typed varinfo to test model type stability with
+    varinfo = DynamicPPL.typed_varinfo(model)
 
-    # Let's make sure that both evaluation and sampling doesn't result in type errors.
-    issuccess, result = DynamicPPL.Experimental.is_suitable_varinfo(
-        sampling_model, varinfo; only_ddpl
+    # Check type stability of evaluation (i.e. DefaultContext)
+    model = DynamicPPL.contextualize(
+        model, DynamicPPL.setleafcontext(model.context, DynamicPPL.DefaultContext())
+    )
+    eval_issuccess, eval_result = DynamicPPL.Experimental.is_suitable_varinfo(
+        model, varinfo; only_ddpl
     )
+    if !eval_issuccess
+        @debug "Evaluation with typed varinfo failed with the following issues:"
+        @debug eval_result
+    end
 
-    if !issuccess
-        # Useful information for debugging.
-        @debug "Evaluaton with typed varinfo failed with the following issues:"
-        @debug result
+    # Check type stability of initialisation (i.e. InitContext)
+    model = DynamicPPL.contextualize(
+        model, DynamicPPL.setleafcontext(model.context, DynamicPPL.InitContext())
+    )
+    init_issuccess, init_result = DynamicPPL.Experimental.is_suitable_varinfo(
+        model, varinfo; only_ddpl
+    )
+    if !init_issuccess
+        @debug "Initialisation with typed varinfo failed with the following issues:"
+        @debug init_result
     end
 
-    # If we didn't fail anywhere, we return the type stable one.
-    return if issuccess
+    # If neither of them failed, we can return the typed varinfo as it's type stable.
+    return if (eval_issuccess && init_issuccess)
         varinfo
     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(sampling_model)
+        DynamicPPL.untyped_varinfo(model)
     end
 end
 

ext/DynamicPPLMCMCChainsExt.jl

@@ -23,7 +23,7 @@ end
 
 function _check_varname_indexing(c::MCMCChains.Chains)
     return DynamicPPL.supports_varname_indexing(c) ||
-           error("Chains do not support indexing using `VarName`s.")
+           error("This `Chains` object does not support indexing using `VarName`s.")
 end
 
 function DynamicPPL.getindex_varname(
@@ -37,6 +37,17 @@ function DynamicPPL.varnames(c::MCMCChains.Chains)
     return keys(c.info.varname_to_symbol)
 end
 
+function chain_sample_to_varname_dict(
+    c::MCMCChains.Chains{Tval}, sample_idx, chain_idx
+) where {Tval}
+    _check_varname_indexing(c)
+    d = Dict{DynamicPPL.VarName,Tval}()
+    for vn in DynamicPPL.varnames(c)
+        d[vn] = DynamicPPL.getindex_varname(c, sample_idx, vn, chain_idx)
+    end
+    return d
+end
+
 """
     predict([rng::AbstractRNG,] model::Model, chain::MCMCChains.Chains; include_all=false)
 
@@ -109,9 +120,15 @@ 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)
-        varinfo = last(DynamicPPL.evaluate_and_sample!!(rng, model, varinfo))
-
+        # Extract values from the chain
+        values_dict = chain_sample_to_varname_dict(parameter_only_chain, sample_idx, chain_idx)
+        # Resample any variables that are not present in `values_dict`
+        _, varinfo = DynamicPPL.init!!(
+            rng,
+            model,
+            varinfo,
+            DynamicPPL.InitFromParams(values_dict, DynamicPPL.InitFromPrior()),
+        )
         vals = DynamicPPL.values_as_in_model(model, false, varinfo)
         varname_vals = mapreduce(
             collect,
@@ -243,13 +260,16 @@ function DynamicPPL.returned(model::DynamicPPL.Model, chain_full::MCMCChains.Cha
     varinfo = DynamicPPL.VarInfo(model)
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
     return map(iters) do (sample_idx, chain_idx)
-        # TODO: Use `fix` once we've addressed https://github.com/TuringLang/DynamicPPL.jl/issues/702.
-        # Update the varinfo with the current sample and make variables not present in `chain`
-        # to be sampled.
-        DynamicPPL.setval_and_resample!(varinfo, chain, sample_idx, chain_idx)
-        # NOTE: Some of the varialbes can be a view into the `varinfo`, so we need to
-        # `deepcopy` the `varinfo` before passing it to the `model`.
-        model(deepcopy(varinfo))
+        # Extract values from the chain
+        values_dict = chain_sample_to_varname_dict(chain, sample_idx, chain_idx)
+        # Resample any variables that are not present in `values_dict`, and
+        # return the model's retval.
+        retval, _ = DynamicPPL.init!!(
+            model,
+            varinfo,
+            DynamicPPL.InitFromParams(values_dict, DynamicPPL.InitFromPrior()),
+        )
+        retval
     end
 end
 

src/DynamicPPL.jl

@@ -206,10 +206,6 @@ include("test_utils.jl")
 include("experimental.jl")
 include("deprecated.jl")
 
-if !isdefined(Base, :get_extension)
-    using Requires
-end
-
 # Better error message if users forget to load JET
 if isdefined(Base.Experimental, :register_error_hint)
     function __init__()

src/extract_priors.jl

@@ -123,7 +123,7 @@ extract_priors(args::Union{Model,AbstractVarInfo}...) =
 function extract_priors(rng::Random.AbstractRNG, model::Model)
     varinfo = VarInfo()
     varinfo = setaccs!!(varinfo, (PriorDistributionAccumulator(),))
-    varinfo = last(evaluate_and_sample!!(rng, model, varinfo))
+    varinfo = last(init!!(rng, model, varinfo))
     return getacc(varinfo, Val(:PriorDistributionAccumulator)).priors
 end
 

src/model.jl

@@ -850,7 +850,7 @@ end
 # ^ Weird Documenter.jl bug means that we have to write the two above separately
 # as it can only detect the `function`-less syntax.
 function (model::Model)(rng::Random.AbstractRNG, varinfo::AbstractVarInfo=VarInfo())
-    return first(evaluate_and_sample!!(rng, model, varinfo))
+    return first(init!!(rng, model, varinfo))
 end
 
 """
@@ -863,32 +863,6 @@ function use_threadsafe_eval(context::AbstractContext, varinfo::AbstractVarInfo)
     return Threads.nthreads() > 1
 end
 
-"""
-    evaluate_and_sample!!([rng::Random.AbstractRNG, ]model::Model, varinfo[, sampler])
-
-Evaluate the `model` with the given `varinfo`, but perform sampling during the
-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 evaluate_and_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 evaluate_and_sample!!(
-    model::Model, varinfo::AbstractVarInfo, sampler::AbstractSampler=SampleFromPrior()
-)
-    return evaluate_and_sample!!(Random.default_rng(), model, varinfo, sampler)
-end
-
 """
     init!!(
         [rng::Random.AbstractRNG,]
@@ -897,12 +871,12 @@ end
         [init_strategy::AbstractInitStrategy=InitFromPrior()]
     )
 
-Evaluate the `model` and replace the values of the model's random variables in
-the given `varinfo` with new values using a specified initialisation strategy.
-If the values in `varinfo` are not already present, they will be added using
-that same strategy.
+Evaluate the `model` and replace the values of the model's random variables
+in the given `varinfo` with new values, using a specified initialisation strategy.
+If the values in `varinfo` are not set, they will be added
+using a specified initialisation strategy.
 
-If `init_strategy` is not provided, defaults to InitFromPrior().
+If `init_strategy` is not provided, defaults to `InitFromPrior()`.
 
 Returns a tuple of the model's return value, plus the updated `varinfo` object.
 """
@@ -1051,11 +1025,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(
-        evaluate_and_sample!!(
-            rng, model, SimpleVarInfo{Float64}(OrderedDict{VarName,Any}())
-        ),
-    )
+    x = last(init!!(rng, model, SimpleVarInfo{Float64}(OrderedDict{VarName,Any}())))
     return values_as(x, T)
 end
 
@@ -1227,25 +1197,8 @@ function Distributions.loglikelihood(model::Model, chain::AbstractMCMC.AbstractC
     end
 end
 
-"""
-    predict([rng::Random.AbstractRNG,] model::Model, chain::AbstractVector{<:AbstractVarInfo})
-
-Generate samples from the posterior predictive distribution by evaluating `model` at each set
-of parameter values provided in `chain`. The number of posterior predictive samples matches
-the length of `chain`. The returned `AbstractVarInfo`s will contain both the posterior parameter values
-and the predicted values.
-"""
-function predict(
-    rng::Random.AbstractRNG, model::Model, chain::AbstractArray{<:AbstractVarInfo}
-)
-    varinfo = DynamicPPL.VarInfo(model)
-    return map(chain) do params_varinfo
-        vi = deepcopy(varinfo)
-        DynamicPPL.setval_and_resample!(vi, values_as(params_varinfo, NamedTuple))
-        model(rng, vi)
-        return vi
-    end
-end
+# Implemented & documented in DynamicPPLMCMCChainsExt
+function predict end
 
 """
     returned(model::Model, parameters::NamedTuple)