Issues with generated_quantities and siblings

[torfjelde]

After recent developments, we now support slightly more complex structures ending up in the chains.

Unfortunately, this means that we're now deviating from the what is expected by much functionality in DynamicPPL, e.g. generated_quantities, see for example #2073.

There are ways we can improve this on the side of DynamicPPL, e.g. by introducing a nested_getindex which can handle scenarios such as

julia> model = DynamicPPL.TestUtils.demo_lkjchol();

julia> vi = VarInfo(model);

julia> DynamicPPL.getindex_nested(vi, @varname(x.L))
2×2 LinearAlgebra.LowerTriangular{Float64, Base.ReshapedArray{Float64, 2, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{}}}:
 1.0          ⋅ 
 0.00433526  0.999991

julia> DynamicPPL.getindex_nested(vi, @varname(x.L[1,1]))
1.0

and similarly a nested_setindex!, which sacrifice performance for flexibility (we'll still keep the existing getindex implementations which will be much faster).

But even something like the above won't be sufficient to fix all our interactions with MCMCChains.Chains objects because at this point everything has been converting into Symbol and so the structure, e.g. is x.L meant to be VarName{:x}(@lens(_.L)) or VarName{Symbol("x.L")) (can easily occur with sub-models using prefixing), has been completely lost.

Unfortunately, due to the reliance on AxisArrays.jl for the underlying storage, we cannot just use VarName as keys for the Chains, as things currently stand.

Sooo we need to do something. There are a few immediate ideas:

1. Make MCMCChains.Chains more flexible.
2. Just add a mapping from varname to symbol to, say, the info field of Chains and define getindex for varname using this.
3. ???

[sunxd3]

@torfjelde #2071 instead of #2073?

[yebai]

Fixed by #2078