Find an edge case

Author: penelopeysm

ext/DynamicPPLMCMCChainsExt.jl

@@ -44,7 +44,7 @@ end
 
 function chain_sample_to_varname_dict(c::MCMCChains.Chains, sample_idx, chain_idx)
     _check_varname_indexing(c)
-    d = Dict{VarName}()
+    d = Dict{DynamicPPL.VarName,Any}()
     for vn in DynamicPPL.varnames(c)
         d[vn] = DynamicPPL.getindex_varname(c, sample_idx, vn, chain_idx)
     end
@@ -271,10 +271,7 @@ function DynamicPPL.returned(model::DynamicPPL.Model, chain_full::MCMCChains.Cha
         # return the model's retval (`first`).
         first(
             DynamicPPL.init!!(
-                rng,
-                model,
-                varinfo,
-                DynamicPPL.ParamsInit(values_dict, DynamicPPL.PriorInit()),
+                model, varinfo, DynamicPPL.ParamsInit(values_dict, DynamicPPL.PriorInit())
             ),
         )
     end

src/contexts/init.jl

@@ -94,6 +94,23 @@ struct ParamsInit{P,S<:AbstractInitStrategy} <: AbstractInitStrategy
     end
 end
 function init(rng::Random.AbstractRNG, vn::VarName, dist::Distribution, p::ParamsInit)
+    # TODO(penelopeysm): Fix this. If anything in p.params _subsumes_ vn,
+    # we don't know how to handle it. This is just another corollary of
+    # https://github.com/TuringLang/DynamicPPL.jl/issues/814
+    # This used to be handled by nested_setindex_maybe, which I'd really like
+    # to get rid of.
+    if p.params isa AbstractDict{<:VarName}
+        strictly_subsumed = filter(
+            vn_in_params -> vn_in_params != vn && subsumes(vn, vn_in_params), keys(p.params)
+        )
+        if !isempty(strictly_subsumed)
+            throw(
+                ArgumentError(
+                    "The given dictionary of parameters contain the following sub-variables of $(vn): $(strictly_subsumed). ParamsInit doesn't know how to deal with this.",
+                ),
+            )
+        end
+    end
     return if hasvalue(p.params, vn)
         x = getvalue(p.params, vn)
         if x === missing
@@ -159,12 +176,15 @@ function tilde_assume(
     y = f(x)
     logjac = logabsdetjac(insert_transformed_value ? link_transform(dist) : identity, x)
     # Add the new value to the VarInfo. `push!!` errors if the value already
-    # exists, hence the need for setindex!!
+    # exists, hence the need for setindex!!.
     if in_varinfo
         vi = setindex!!(vi, y, vn)
     else
         vi = push!!(vi, vn, y, dist)
     end
+    # Neither of these set the `trans` flag so we have to do it manually if
+    # necessary.
+    insert_transformed_value && settrans!!(vi, true, vn)
     # `accumulate_assume!!` wants untransformed values as the second argument.
     vi = accumulate_assume!!(vi, x, -logjac, vn, dist)
     # We always return the untransformed value here, as that will determine

src/simple_varinfo.jl

@@ -463,7 +463,6 @@ end
 
 # Context implementations
 
-# NOTE: We don't implement `settrans!!(vi, trans, vn)`.
 function settrans!!(vi::SimpleVarInfo, trans)
     return settrans!!(vi, trans ? DynamicTransformation() : NoTransformation())
 end
@@ -473,6 +472,9 @@ end
 function settrans!!(vi::ThreadSafeVarInfo{<:SimpleVarInfo}, trans)
     return Accessors.@set vi.varinfo = settrans!!(vi.varinfo, trans)
 end
+function settrans!!(::SimpleOrThreadSafeSimple, trans::Bool, vn::VarName)
+    @info "Attempting to call `settrans!!` on a `SimpleVarInfo` for a specific variable `$vn`; this will be ignored"
+end
 
 istrans(vi::SimpleVarInfo) = !(vi.transformation isa NoTransformation)
 istrans(vi::SimpleVarInfo, ::VarName) = istrans(vi)

test/test_util.jl

@@ -81,8 +81,10 @@ function make_chain_from_prior(rng::Random.AbstractRNG, model::Model, n_iters::I
     varnames = collect(varnames)
     # Construct matrix of values
     vals = [get(dict, vn, missing) for dict in dicts, vn in varnames]
+    # Construct mapping of varnames to symbols
+    vns_to_syms = Dict{VarName,Symbol}(zip(varnames, Symbol.(varnames)))
     # Construct and return the Chains object
-    return Chains(vals, varnames)
+    return Chains(vals, varnames; info=(varname_to_symbol=vns_to_syms,))
 end
 function make_chain_from_prior(model::Model, n_iters::Int)
     return make_chain_from_prior(Random.default_rng(), model, n_iters)