Indexing into `SimpleVarInfo`

[ParadaCarleton]

julia> @model function demo()
              μ ~ Normal()
              x = Vector{Float64}(undef, 2)
              for i in eachindex(x)
                  x[i] ~ Normal()
              end
              return x
          end
demo (generic function with 2 methods)

julia> m = demo()
Model{typeof(demo), (), (), (), Tuple{}, Tuple{}, DefaultContext}(:demo, demo, NamedTuple(), NamedTuple(), DefaultContext())

julia> vi = SimpleVarInfo(VarInfo(m))

julia> logjoint(m, vi)
ERROR: type NamedTuple has no field x
Stacktrace:
  [1] getproperty
    @ ./Base.jl:42 [inlined]
  [2] get(obj::NamedTuple{(:μ, Symbol("x[1]"), Symbol("x[2]")), Tuple{Float64, Float64, Float64}}, l::Setfield.PropertyLens{:x})
    @ Setfield ~/.julia/packages/Setfield/9IaGB/src/lens.jl:117
  [3] get(obj::NamedTuple{(:μ, Symbol("x[1]"), Symbol("x[2]")), Tuple{Float64, Float64, Float64}}, l::Setfield.ComposedLens{Setfield.PropertyLens{:x}, Setfield.IndexLens{Tuple{Int64}}})
    @ Setfield ~/.julia/packages/Setfield/9IaGB/src/lens.jl:182
  [4] get(obj::NamedTuple{(:μ, Symbol("x[1]"), Symbol("x[2]")), Tuple{Float64, Float64, Float64}}, vn::VarName{:x, Setfield.IndexLens{Tuple{Int64}}})
    @ AbstractPPL ~/.julia/packages/AbstractPPL/EkJZe/src/varname.jl:131
  [5] getindex(vi::SimpleVarInfo{NamedTuple{(:μ, Symbol("x[1]"), Symbol("x[2]")), Tuple{Float64, Float64, Float64}}, Float64}, vn::VarName{:x, Setfield.IndexLens{Tuple{Int64}}})
    @ DynamicPPL ~/.julia/dev/DynamicPPL.jl/src/simple_varinfo.jl:186
  [6] getindex(vi::DynamicPPL.ThreadSafeVarInfo{SimpleVarInfo{NamedTuple{(:μ, Symbol("x[1]"), Symbol("x[2]")), Tuple{Float64, Float64, Float64}}, Float64}, Vector{Base.RefValue{Float64}}}, vn::VarName{:x, Setfield.IndexLens{Tuple{Int64}}})
    @ DynamicPPL ~/.julia/dev/DynamicPPL.jl/src/threadsafe.jl:65
  [7] assume(dist::Normal{Float64}, vn::VarName{:x, Setfield.IndexLens{Tuple{Int64}}}, vi::DynamicPPL.ThreadSafeVarInfo{SimpleVarInfo{NamedTuple{(:μ, Symbol("x[1]"), Symbol("x[2]")), Tuple{Float64, Float64, Float64}}, Float64}, Vector{Base.RefValue{Float64}}})
    @ DynamicPPL ~/.julia/dev/DynamicPPL.jl/src/simple_varinfo.jl:341
  [8] tilde_assume(#unused#::DynamicPPL.IsLeaf, context::DefaultContext, right::Normal{Float64}, vn::VarName{:x, Setfield.IndexLens{Tuple{Int64}}}, vi::DynamicPPL.ThreadSafeVarInfo{SimpleVarInfo{NamedTuple{(:μ, Symbol("x[1]"), Symbol("x[2]")), Tuple{Float64, Float64, Float64}}, Float64}, Vector{Base.RefValue{Float64}}})
    @ DynamicPPL ~/.julia/dev/DynamicPPL.jl/src/context_implementations.jl:39
  [9] tilde_assume(::DefaultContext, ::Normal{Float64}, ::VarName{:x, Setfield.IndexLens{Tuple{Int64}}}, ::DynamicPPL.ThreadSafeVarInfo{SimpleVarInfo{NamedTuple{(:μ, Symbol("x[1]"), Symbol("x[2]")), Tuple{Float64, Float64, Float64}}, Float64}, Vector{Base.RefValue{Float64}}})
    @ DynamicPPL ~/.julia/dev/DynamicPPL.jl/src/context_implementations.jl:36
 [10] tilde_assume!!(context::DefaultContext, right::Normal{Float64}, vn::VarName{:x, Setfield.IndexLens{Tuple{Int64}}}, vi::DynamicPPL.ThreadSafeVarInfo{SimpleVarInfo{NamedTuple{(:μ, Symbol("x[1]"), Symbol("x[2]")), Tuple{Float64, Float64, Float64}}, Float64}, Vector{Base.RefValue{Float64}}})
    @ DynamicPPL ~/.julia/dev/DynamicPPL.jl/src/context_implementations.jl:117
 [11] demo
    @ ./REPL[3]:4 [inlined]
 [12] macro expansion
    @ ~/.julia/dev/DynamicPPL.jl/src/model.jl:489 [inlined]
 [13] _evaluate!!
    @ ~/.julia/dev/DynamicPPL.jl/src/model.jl:472 [inlined]
 [14] evaluate_threadsafe!!(model::Model{typeof(demo), (), (), (), Tuple{}, Tuple{}, DefaultContext}, varinfo::SimpleVarInfo{NamedTuple{(:μ, Symbol("x[1]"), Symbol("x[2]")), Tuple{Float64, Float64, Float64}}, Float64}, context::DefaultContext)
    @ DynamicPPL ~/.julia/dev/DynamicPPL.jl/src/model.jl:463
 [15] evaluate!!
    @ ~/.julia/dev/DynamicPPL.jl/src/model.jl:402 [inlined]
 [16] logjoint(model::Model{typeof(demo), (), (), (), Tuple{}, Tuple{}, DefaultContext}, varinfo::SimpleVarInfo{NamedTuple{(:μ, Symbol("x[1]"), Symbol("x[2]")), Tuple{Float64, Float64, Float64}}, Float64})
    @ DynamicPPL ~/.julia/dev/DynamicPPL.jl/src/model.jl:522
 [17] top-level scope
    @ REPL[18]:1

julia> vi
SimpleVarInfo((μ = 0.35792840316913777, x[1] = -0.05997503801481726, x[2] = -2.2264843907029834), -5.30128684412608)

Looks like indexing with :x into SimpleVarInfo doesn't work, so logjoint fails.

Given SimpleVarInfo is a named tuple that can store arbitrary types, is there anything keeping us from storing arrays in it?

[yebai]

@ParadaCarleton It is a known limitation of SimpleVarInfo. For now, we have to provide arrays for SimpleVarInfo, e.g.:

vi = SimpleVarInfo(((μ = 0.35792840316913777, x = rand(2)))

logjoint(m, vi)

In vi = SimpleVarInfo(VarInfo(m)), the SimpleVarInfo constructor function treats x[1] and x[2] as two scalars with special names "x[1]" and "x[2]", instead of treating x as a vector (x does not exist in fact). It is possible to support it but intentionally left as a future improvement for now. In order to support it, we need to implement a convert function from ("x[1]" = 0, "x[2]"=1) to x = [0 1], which is not trivial because we need to infer the shape of x.

[ParadaCarleton]

Makes sense. What about logjoint failing? It should probably work with array-valued variables.

[torfjelde]

Have a look at the SimpleVarInfo you created and you'll find that it doesn't have x in there, it only has var"x[1]" and var"x[2]" . This stuff is the limitation of SimpleVarInfo; I recommend looking at the docstring, as it tries to outline when it fails, and if it's not descriptive enough we should improve that 👍

If you want something that just works on everything, try SimpleVarInfo using a Dict.

[ParadaCarleton]

Have a look at the SimpleVarInfo you created and you'll find that it doesn't have x in there, it only has var"x[1]" and var"x[2]" . This stuff is the limitation of SimpleVarInfo; I recommend looking at the docstring, as it tries to outline when it fails, and if it's not descriptive enough we should improve that +1

If you want something that just works on everything, try SimpleVarInfo using a Dict.

Makes sense. What about logjoint not working for SimpleVarInfo like this?

[torfjelde]

Makes sense. What about logjoint not working for SimpleVarInfo like this?

That's what I was trying to answer 😅

So getindex for SimpleVarInfo assumes the variable name, in this case x, is present in it's values field, that is varinfo.values. This makes it real fast, but it comes at the disadvantage that it won't work when you have (var"x[1]" = ..., var"x[2]" = ...) (this quickly becomes infeasible as the length of x grows, hence why we didn't even bother supporting it). So when you try to evaluate the logjoint, we hit x[1] ~ Normal() which, at some point in the call stack, results in getindex(varinfo, @varname(x[1])). This will fail when because it can't find x in varinfo.values.

To make it work for this model, you have to either:

1. Manually construct the varinfo, i.e. SimpleVarInfo((x = randn(n), )).
2. Construct a SimpleVarInfo with Dict instead which supports more general indexing, e.g. if you do SimpleVarInfo(Dict(@varname(x) => [1.0, 2.0])) then getindex(varinfo, @varname(x)) and getindex(varinfo, @varname(x[1])) both work as expected, and SimpleVarInfo(Dict(@varname(x[1]) => 1.0, @varname(x[2]) => 2.0)) will work with getindex(varinfo, @varname(x[1])) but not getindex(varinfo, @varname(x)). Essentially, the keys present in varinfo.values must subsume (I recommend the doc in AdvancedPPL.jl for more on what I mean by this) the VarName used in the call to getindex.

[yebai]

@ParadaCarleton Please feel free to re-open this issue if you got any further questions.