Fix a bunch of tests

Author: penelopeysm

ext/DynamicPPLMCMCChainsExt.jl

@@ -42,6 +42,15 @@ 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)
+    _check_varname_indexing(c)
+    d = Dict{VarName}()
+    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)
 
@@ -114,9 +123,19 @@ 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 = DynamicPPL.chain_sample_to_varname_dict(
+            parameter_only_chain, sample_idx, chain_idx
+        )
+        # Resample any variables that are not present in `values_dict`
+        _, varinfo = last(
+            DynamicPPL.init!!(
+                rng,
+                model,
+                varinfo,
+                DynamicPPL.ParamsInit(values_dict, DynamicPPL.PriorInit()),
+            ),
+        )
         vals = DynamicPPL.values_as_in_model(model, false, varinfo)
         varname_vals = mapreduce(
             collect,
@@ -248,13 +267,20 @@ 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 = DynamicPPL.chain_sample_to_varname_dict(
+            parameter_only_chain, sample_idx, chain_idx
+        )
+        # Resample any variables that are not present in `values_dict`, and
+        # return the model's retval (`first`).
+        first(
+            DynamicPPL.init!!(
+                rng,
+                model,
+                varinfo,
+                DynamicPPL.ParamsInit(values_dict, DynamicPPL.PriorInit()),
+            ),
+        )
     end
 end
 

src/debug_utils.jl

@@ -438,7 +438,9 @@ function check_model_and_trace(
     kwargs...,
 )
     # Execute the model with the debug context.
-    new_context = setleafcontext(model.context, InitContext(rng, Prior()))
+    new_context = DynamicPPL.setleafcontext(
+        model.context, DynamicPPL.InitContext(rng, DynamicPPL.PriorInit())
+    )
     debug_context = DebugContext(new_context; error_on_failure=error_on_failure, kwargs...)
     debug_model = DynamicPPL.contextualize(model, debug_context)
 

src/model.jl

@@ -1165,8 +1165,8 @@ function predict(
     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)
+        values_nt = values_as(params_varinfo, NamedTuple)
+        _, vi = DynamicPPL.init!!(rng, model, vi, ParamsInit(values_nt, PriorInit()))
         return vi
     end
 end

src/test_utils/contexts.jl

@@ -29,21 +29,45 @@ function test_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Mod
     node_trait = DynamicPPL.NodeTrait(context)
     # Throw error immediately if it it's missing a `NodeTrait` implementation.
     node_trait isa Union{DynamicPPL.IsLeaf,DynamicPPL.IsParent} ||
-        throw(ValueError("Invalid NodeTrait: $node_trait"))
+        error("Invalid NodeTrait: $node_trait")
 
-    # To see change, let's make sure we're using a different leaf context than the current.
-    leafcontext_new = if DynamicPPL.leafcontext(context) isa DefaultContext
-        DynamicPPL.DynamicTransformationContext{false}()
+    if node_trait isa DynamicPPL.IsLeaf
+        test_leaf_context(context, model)
     else
-        DefaultContext()
+        test_parent_context(context, model)
     end
-    @test DynamicPPL.leafcontext(DynamicPPL.setleafcontext(context, leafcontext_new)) ==
-        leafcontext_new
+end
+
+function test_leaf_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Model)
+    @test DynamicPPL.NodeTrait(context) isa DynamicPPL.IsLeaf
+
+    # Note that for a leaf context we can't assume that it will work with an
+    # empty VarInfo. Thus we only test evaluation (i.e., assuming that the
+    # varinfo already contains all necessary variables).
+    @testset "evaluation" begin
+        # Generate a new filled untyped varinfo
+        _, untyped_vi = DynamicPPL.init!!(model, DynamicPPL.VarInfo())
+        typed_vi = DynamicPPL.typed_varinfo(untyped_vi)
+        new_model = contextualize(model, context)
+        for vi in [untyped_vi, typed_vi]
+            _, vi = DynamicPPL.evaluate!!(new_model, vi)
+            @test vi isa DynamicPPL.VarInfo
+        end
+    end
+end
+
+function test_parent_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Model)
+    @test DynamicPPL.NodeTrait(context) isa DynamicPPL.IsParent
 
-    # The interface methods.
-    if node_trait isa DynamicPPL.IsParent
-        # `childcontext` and `setchildcontext`
-        # With new child context
+    @testset "{set,}{leaf,child}context" begin
+        # Ensure we're using a different leaf context than the current.
+        leafcontext_new = if DynamicPPL.leafcontext(context) isa DefaultContext
+            DynamicPPL.DynamicTransformationContext{false}()
+        else
+            DefaultContext()
+        end
+        @test DynamicPPL.leafcontext(DynamicPPL.setleafcontext(context, leafcontext_new)) ==
+            leafcontext_new
         childcontext_new = TestParentContext()
         @test DynamicPPL.childcontext(
             DynamicPPL.setchildcontext(context, childcontext_new)
@@ -56,15 +80,15 @@ function test_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Mod
             leafcontext_new
     end
 
-    # Make sure that the we can evaluate the model with the context (i.e. that none of the tilde-functions are incorrectly overloaded).
-    # NOTE(torfjelde): Need to sample with the untyped varinfo _using_ the context, since the
-    # context might alter which variables are present, their names, etc., e.g. `PrefixContext`.
-    # TODO(torfjelde): Make the `varinfo` used for testing a kwarg once it makes sense for other varinfos.
-    # Untyped varinfo.
-    varinfo_untyped = DynamicPPL.VarInfo()
-    new_model = contextualize(model, context)
-    @test DynamicPPL.evaluate!!(new_model, varinfo_untyped) isa Any
-    # Typed varinfo.
-    varinfo_typed = DynamicPPL.typed_varinfo(varinfo_untyped)
-    @test DynamicPPL.evaluate!!(new_model, varinfo_typed) isa Any
+    @testset "initialisation and evaluation" begin
+        new_model = contextualize(model, context)
+        for vi in [DynamicPPL.VarInfo(), DynamicPPL.typed_varinfo(DynamicPPL.VarInfo())]
+            # Initialisation
+            _, vi = DynamicPPL.init!!(new_model, DynamicPPL.VarInfo())
+            @test vi isa DynamicPPL.VarInfo
+            # Evaluation
+            _, vi = DynamicPPL.evaluate!!(new_model, vi)
+            @test vi isa DynamicPPL.VarInfo
+        end
+    end
 end

src/varinfo.jl

@@ -2045,113 +2045,6 @@ function _setval_kernel!(vi::VarInfoOrThreadSafeVarInfo, vn::VarName, values, ke
     return indices
 end
 
-"""
-    setval_and_resample!(vi::VarInfo, x)
-    setval_and_resample!(vi::VarInfo, values, keys)
-    setval_and_resample!(vi::VarInfo, chains::AbstractChains, sample_idx, chain_idx)
-
-Set the values in `vi` to the provided values and those which are not present
-in `x` or `chains` to *be* resampled.
-
-Note that this does *not* resample the values not provided! It will call
-`setflag!(vi, vn, "del")` for variables `vn` for which no values are provided, which means
-that the next time we call `model(vi)` these variables will be resampled.
-
-## Note
-- This suffers from the same limitations as [`setval!`](@ref). See `setval!` for more info.
-
-## Example
-```jldoctest
-julia> using DynamicPPL, Distributions, StableRNGs
-
-julia> @model function demo(x)
-           m ~ Normal()
-           for i in eachindex(x)
-               x[i] ~ Normal(m, 1)
-           end
-       end;
-
-julia> rng = StableRNG(42);
-
-julia> m = demo([missing]);
-
-julia> var_info = DynamicPPL.VarInfo(rng, m);
-       # Checking the setting of "del" flags only makes sense for VarInfo{<:Metadata}. For VarInfo{<:VarNamedVector} the flag is effectively always set.
-
-julia> var_info[@varname(m)]
--0.6702516921145671
-
-julia> var_info[@varname(x[1])]
--0.22312984965118443
-
-julia> DynamicPPL.setval_and_resample!(var_info, (m = 100.0, )); # set `m` and ready `x[1]` for resampling
-
-julia> var_info[@varname(m)] # [✓] changed
-100.0
-
-julia> var_info[@varname(x[1])] # [✓] unchanged
--0.22312984965118443
-
-julia> m(rng, var_info); # sample `x[1]` conditioned on `m = 100.0`
-
-julia> var_info[@varname(m)] # [✓] unchanged
-100.0
-
-julia> var_info[@varname(x[1])] # [✓] changed
-101.37363069798343
-```
-
-## See also
-- [`setval!`](@ref)
-"""
-function setval_and_resample!(vi::VarInfoOrThreadSafeVarInfo, x)
-    return setval_and_resample!(vi, values(x), keys(x))
-end
-function setval_and_resample!(vi::VarInfoOrThreadSafeVarInfo, values, keys)
-    return _apply!(_setval_and_resample_kernel!, vi, values, keys)
-end
-function setval_and_resample!(
-    vi::VarInfoOrThreadSafeVarInfo, chains::AbstractChains, sample_idx::Int, chain_idx::Int
-)
-    if supports_varname_indexing(chains)
-        # First we need to set every variable to be resampled.
-        for vn in keys(vi)
-            set_flag!(vi, vn, "del")
-        end
-        # Then we set the variables in `varinfo` from `chain`.
-        for vn in varnames(chains)
-            vn_updated = nested_setindex_maybe!(
-                vi, getindex_varname(chains, sample_idx, vn, chain_idx), vn
-            )
-
-            # Unset the `del` flag if we found something.
-            if vn_updated !== nothing
-                # NOTE: This will be triggered even if only a subset of a variable has been set!
-                unset_flag!(vi, vn_updated, "del")
-            end
-        end
-    else
-        setval_and_resample!(vi, chains.value[sample_idx, :, chain_idx], keys(chains))
-    end
-end
-
-function _setval_and_resample_kernel!(
-    vi::VarInfoOrThreadSafeVarInfo, vn::VarName, values, keys
-)
-    indices = findall(Base.Fix1(subsumes_string, string(vn)), keys)
-    if !isempty(indices)
-        val = reduce(vcat, values[indices])
-        setval!(vi, val, vn)
-        settrans!!(vi, false, vn)
-    else
-        # Ensures that we'll resample the variable corresponding to `vn` if we run
-        # the model on `vi` again.
-        set_flag!(vi, vn, "del")
-    end
-
-    return indices
-end
-
 values_as(vi::VarInfo) = vi.metadata
 values_as(vi::VarInfo, ::Type{Vector}) = copy(getindex_internal(vi, Colon()))
 function values_as(vi::UntypedVarInfo, ::Type{NamedTuple})

test/compiler.jl

@@ -194,7 +194,7 @@ module Issue537 end
         @test getlogjoint(varinfo) == lp
         @test varinfo_ isa AbstractVarInfo
         @test model_.f === model.f
-        @test model_.context isa InitContext
+        @test model_.context isa DynamicPPL.InitContext
         @test model_.context.rng isa Random.AbstractRNG
 
         # disable warnings
@@ -595,13 +595,13 @@ module Issue537 end
         # an attempt at a `NamedTuple` of the form `(x = 1, __varinfo__)`.
         @model empty_model() = return x = 1
         empty_vi = VarInfo()
-        retval_and_vi = DynamicPPL.evaluate_and_sample!!(empty_model(), empty_vi)
+        retval_and_vi = DynamicPPL.init!!(empty_model(), empty_vi)
         @test retval_and_vi isa Tuple{Int,typeof(empty_vi)}
 
         # Even if the return-value is `AbstractVarInfo`, we should return
         # a `Tuple` with `AbstractVarInfo` in the second component too.
         @model demo() = return __varinfo__
-        retval, svi = DynamicPPL.evaluate_and_sample!!(demo(), SimpleVarInfo())
+        retval, svi = DynamicPPL.init!!(demo(), SimpleVarInfo())
         @test svi == SimpleVarInfo()
         if Threads.nthreads() > 1
             @test retval isa DynamicPPL.ThreadSafeVarInfo{<:SimpleVarInfo}
@@ -617,11 +617,11 @@ module Issue537 end
             f(x) = return x^2
             return f(1.0)
         end
-        retval, svi = DynamicPPL.evaluate_and_sample!!(demo(), SimpleVarInfo())
+        retval, svi = DynamicPPL.init!!(demo(), SimpleVarInfo())
         @test retval isa Float64
 
         @model demo() = x ~ Normal()
-        retval, svi = DynamicPPL.evaluate_and_sample!!(demo(), SimpleVarInfo())
+        retval, svi = DynamicPPL.init!!(demo(), SimpleVarInfo())
 
         # Return-value when using `to_submodel`
         @model inner() = x ~ Normal()

test/ext/DynamicPPLJETExt.jl

@@ -70,10 +70,6 @@
             )
             JET.test_call(f_eval, argtypes_eval)
 
-            f_sample, argtypes_sample = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
-                init_model, varinfo
-            )
-            JET.test_call(f_sample, argtypes_sample)
             # For our demo models, they should all result in typed.
             is_typed = varinfo isa DynamicPPL.NTVarInfo
             @test is_typed