Support for variables with changing size with dot_tilde

[mhauru]

Currently a variable that expands in size between iterations crashes `dot_tilde:

MWE:

using Random, DynamicPPL, Distributions

Random.seed!(23)

@model function dynamic_model_with_dot_tilde(num_zs=10)
    z = Vector(undef, num_zs)
    z .~ Exponential(1.0)
    num_ms = Int(round(sum(z)))
    @show num_ms
    m = Vector(undef, num_ms)
    return m .~ Normal(1.0, 1.0)
end
m = dynamic_model_with_dot_tilde()

vi = VarInfo(m)
for k in keys(vi)
    set_flag!(vi, k, "del")
end
m(vi)

Output:

num_ms = 8
num_ms = 11
ERROR: KeyError: key m[9] not found
Stacktrace:
  [1] getindex
    @ ./dict.jl:498 [inlined]
  [2] getidx
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/varinfo.jl:605 [inlined]
  [3] is_flagged
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/varinfo.jl:1940 [inlined]
  [4] istrans
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/varinfo.jl:1073 [inlined]
  [5] istrans
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/varinfo.jl:1072 [inlined]
  [6] from_maybe_linked_internal_transform
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/abstract_varinfo.jl:828 [inlined]
  [7] to_maybe_linked_internal_transform
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/abstract_varinfo.jl:881 [inlined]
  [8] get_and_set_val!(rng::TaskLocalRNG, vi::TypedVarInfo{@NamedTuple{z::DynamicPPL.Metadata{…}, m::DynamicPPL.Metadata{…}}, Float64}, vns::Vector{VarName{:m, Accessors.IndexLens{…}}}, dists::Normal{Float64}, spl::SampleFromPrior)
    @ DynamicPPL ~/.julia/packages/DynamicPPL/L81Uf/src/context_implementations.jl:574
  [9] dot_assume(rng::TaskLocalRNG, spl::SampleFromPrior, dists::Normal{Float64}, vns::Vector{VarName{:m, Accessors.IndexLens{…}}}, var::Vector{Any}, vi::TypedVarInfo{@NamedTuple{z::DynamicPPL.Metadata{…}, m::DynamicPPL.Metadata{…}}, Float64})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/L81Uf/src/context_implementations.jl:489
 [10] dot_tilde_assume
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/context_implementations.jl:327 [inlined]
 [11] dot_tilde_assume
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/context_implementations.jl:297 [inlined]
 [12] dot_tilde_assume!!
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/context_implementations.jl:423 [inlined]
 [13] macro expansion
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/compiler.jl:579 [inlined]
 [14] dynamic_model_with_dot_tilde(__model__::Model{…}, __varinfo__::TypedVarInfo{…}, __context__::SamplingContext{…}, num_zs::Int64)
    @ Main.MWE ./REPL[29]:12
 [15] _evaluate!!
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/model.jl:975 [inlined]
 [16] evaluate_threadunsafe!!
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/model.jl:948 [inlined]
 [17] evaluate!!
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/model.jl:896 [inlined]
 [18] evaluate!! (repeats 2 times)
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/model.jl:907 [inlined]
 [19] evaluate!!
    @ ~/.julia/packages/DynamicPPL/L81Uf/src/model.jl:917 [inlined]
 [20] (::Model{typeof(Main.MWE.dynamic_model_with_dot_tilde), (:num_zs,), (), (), Tuple{Int64}, Tuple{}, DefaultContext})(args::TypedVarInfo{@NamedTuple{z::DynamicPPL.Metadata{…}, m::DynamicPPL.Metadata{…}}, Float64})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/L81Uf/src/model.jl:867
 [21] top-level scope
    @ REPL[29]:20
Some type information was truncated. Use `show(err)` to see complete types.

[torfjelde]

This is caused by lots of checks in get_and_set_val! (which is a very ugly function itself) performing some checks based on vns[1] and others based on vns[i].

We should just clean up that impl and this will be good 👍