InitContext, part 4 - Use init!! to replace evaluate_and_sample!!, predict, returned, and initialize_values

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"

ext/DynamicPPLMCMCChainsExt.jl

@@ -37,9 +37,11 @@ function DynamicPPL.varnames(c::MCMCChains.Chains)
     return keys(c.info.varname_to_symbol)
 end
 
-function chain_sample_to_varname_dict(c::MCMCChains.Chains, sample_idx, chain_idx)
+function chain_sample_to_varname_dict(
+    c::MCMCChains.Chains{Tval}, sample_idx, chain_idx
+) where {Tval}
     _check_varname_indexing(c)
-    d = Dict{DynamicPPL.VarName,Any}()
+    d = Dict{DynamicPPL.VarName,Tval}()
     for vn in DynamicPPL.varnames(c)
         d[vn] = DynamicPPL.getindex_varname(c, sample_idx, vn, chain_idx)
     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/model.jl

@@ -873,10 +873,10 @@ end
 
 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.
+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.
 """

src/sampler.jl

@@ -41,7 +41,7 @@ Generic sampler type for inference algorithms of type `T` in DynamicPPL.
 provided that supports resuming sampling from a previous state and setting initial
 parameter values. It requires to overload [`loadstate`](@ref) and [`initialstep`](@ref)
 for loading previous states and actually performing the initial sampling step,
-respectively. Additionally, sometimes one might want to implement [`init_strategy`](@ref)
+respectively. Additionally, sometimes one might want to implement an [`init_strategy`](@ref)
 that specifies how the initial parameter values are sampled if they are not provided.
 By default, values are sampled from the prior.
 """
@@ -68,7 +68,7 @@ end
 
 Return a default varinfo object for the given `model` and `sampler`.
 
-The default method for this returns an empty NTVarInfo (i.e. 'typed varinfo').
+The default method for this returns a NTVarInfo (i.e. 'typed varinfo').
 
 # Arguments
 - `rng::Random.AbstractRNG`: Random number generator.
@@ -78,12 +78,24 @@ The default method for this returns an empty NTVarInfo (i.e. 'typed varinfo').
 # Returns
 - `AbstractVarInfo`: Default varinfo object for the given `model` and `sampler`.
 """
-function default_varinfo(::Random.AbstractRNG, ::Model, ::AbstractSampler)
-    # Note that variable values are unconditionally initialized later, so no
-    # point putting them in now.
-    return typed_varinfo(VarInfo())
+function default_varinfo(rng::Random.AbstractRNG, model::Model, ::AbstractSampler)
+    # Note that in `AbstractMCMC.step`, the values in the varinfo returned here are
+    # immediately overwritten by a subsequent call to `init!!`. The reason why we
+    # _do_ create a varinfo with parameters here (as opposed to simply returning
+    # an empty `typed_varinfo(VarInfo())`) is to avoid issues where pushing to an empty
+    # typed VarInfo would fail. This can happen if two VarNames have different types
+    # but share the same symbol (e.g. `x.a` and `x.b`).
+    return typed_varinfo(VarInfo(rng, model))
 end
 
+"""
+    init_strategy(sampler)
+
+Define the initialisation strategy used for generating initial values when
+sampling with `sampler`. Defaults to `InitFromPrior()`, but can be overridden.
+"""
+init_strategy(::Sampler) = InitFromPrior()
+
 function AbstractMCMC.sample(
     rng::Random.AbstractRNG,
     model::Model,
@@ -99,13 +111,22 @@ function AbstractMCMC.sample(
     )
 end
 
-"""
-    init_strategy(sampler)
-
-Define the initialisation strategy used for generating initial values when
-sampling with `sampler`. Defaults to `InitFromPrior()`, but can be overridden.
-"""
-init_strategy(::Sampler) = InitFromPrior()
+function AbstractMCMC.sample(
+    rng::Random.AbstractRNG,
+    model::Model,
+    sampler::Sampler,
+    parallel::AbstractMCMC.AbstractMCMCEnsemble,
+    N::Integer,
+    nchains::Integer;
+    chain_type=default_chain_type(sampler),
+    resume_from=nothing,
+    initial_state=loadstate(resume_from),
+    kwargs...,
+)
+    return AbstractMCMC.mcmcsample(
+        rng, model, sampler, parallel, N, nchains; chain_type, initial_state, kwargs...
+    )
+end
 
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
@@ -114,8 +135,8 @@ function AbstractMCMC.step(
     initial_params::AbstractInitStrategy=init_strategy(spl),
     kwargs...,
 )
-    # Generate the default varinfo (usually this just makes an empty VarInfo
-    # with NamedTuple of Metadata).
+    # Generate the default varinfo. Note that any parameters inside this varinfo
+    # will be immediately overwritten by the next call to `init!!`.
     vi = default_varinfo(rng, model, spl)
 
     # Fill it with initial parameters. Note that, if `InitFromParams` is used, the

src/test_utils/contexts.jl

@@ -59,7 +59,7 @@ end
 function test_parent_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Model)
     @test DynamicPPL.NodeTrait(context) isa DynamicPPL.IsParent
 
-    @testset "{set,}{leaf,child}context" begin
+    @testset "get/set leaf and child contexts" begin
         # Ensure we're using a different leaf context than the current.
         leafcontext_new = if DynamicPPL.leafcontext(context) isa DefaultContext
             DynamicPPL.DynamicTransformationContext{false}()

test/compiler.jl

@@ -193,6 +193,9 @@ module Issue537 end
         varinfo = VarInfo(model)
         @test getlogjoint(varinfo) == lp
         @test varinfo_ isa AbstractVarInfo
+        # During the model evaluation, its leaf context is changed to an InitContext, so
+        # `model_` is not going to be equal to `model`. We can still check equality of `f`
+        # though.
         @test model_.f === model.f
         @test model_.context isa DynamicPPL.InitContext
         @test model_.context.rng isa Random.AbstractRNG

test/sampler.jl

@@ -1,4 +1,17 @@
 @testset "sampler.jl" begin
+    @testset "varnames with same symbol but different type" begin
+        struct S <: AbstractMCMC.AbstractSampler end
+        DynamicPPL.initialstep(rng, model, ::DynamicPPL.Sampler{S}, vi; kwargs...) = vi
+        @model function g()
+            y = (; a=1, b=2)
+            y.a ~ Normal()
+            return y.b ~ Normal()
+        end
+        model = g()
+        spl = DynamicPPL.Sampler(S())
+        @test AbstractMCMC.step(Xoshiro(468), g(), spl) isa Any
+    end
+
     @testset "initial_state and resume_from kwargs" begin
         # Model is unused, but has to be a DynamicPPL.Model otherwise we won't hit our
         # overloaded method.