Return an AbstractChains object from appliying pointwise_logdensities to Abstractchains?

[bgctw]

Currently, pointwise_logdensities(chains::AbstractChains) returns an OrderedDict of Matrices all of the shape (sample x chain).

I find it more convenient to work with the results of the pointwise functions applied to AbstractChains as an AbstractChains again, rather than the OrderedDict{String, Matrix}.
What is a good place to support this conversion? A function in the MCMCChains extension such as:

function as_chains(lds_pointwise)
     Chains(stack(values(lds_pointwise); dims=2), collect(keys(lds_pointwise)))
end
chn = as_chains(logjoints_pointwise); # from @testset "pointwise_logdensities chain"
names(chn)
get(chn, :x)[1] == logjoints_pointwise["x"]

One could even think of letting the pointwise_logdensities(..., ::AbstractChains) routinely return a Chains object.
This could be achieved by

• renaming the pointwise_logdensities(..., ::AbstractChains) to pointwise_logdensities_dict,
• implementing it in the DynamicPPLMCMCChains extension by converting the result of pointwise_logdensities_dict.

What are opinions about this suggestion?

[penelopeysm]

Couple of thoughts:

• From a practical perspective it makes a lot of sense to use a chain type.
• But I don't fully like that in one Chains, chn[:y] could represent a value whereas in another chn[:y] could represent the logp associated with that value, and from a software engineering perspective, I don't want DynamicPPL to return a chain that could be used wrongly downstream.

Based on the balance between this, I think there are a couple of options we could explore:

(1) keep the existing behaviour of pointwise_logdensities, but implement the as_chains function that you have somewhere. I think putting it in MCMCChains.jl directly would be the way to go. This would be a cheap fix

(2) I think the following is a better fix, although as I'll explain, it's not a super easy thing to fix. MCMCChains.Chains actually has this notion of "sections", which is not super enforced:

julia> @model function f(b)
           a ~ Normal()
           b ~ Normal(a)
       end
f (generic function with 2 methods)

julia> model = f(1.0)
Model{typeof(f), (:b,), (), (), Tuple{Float64}, Tuple{}, DefaultContext}(f, (b = 1.0,), NamedTuple(), DefaultContext())

julia> chn = sample(model, NUTS(), 100; progress=false);
┌ Info: Found initial step size
└   ϵ = 1.6

julia> get(chn; section=:internals)[:lp]
2-dimensional AxisArray{Float64,2,...} with axes:
    :iter, 51:1:150
    :chain, 1:1
And data, a 100×1 Matrix{Float64}:
 -4.413545355742071
 -2.2827550349443486
[...]

julia> get(chn; section=:parameters)[:lp]
ERROR: type NamedTuple has no field lp
[...]

julia> get(chn; section=:parameters)[:a]
2-dimensional AxisArray{Float64,2,...} with axes:
    :iter, 51:1:150
    :chain, 1:1
And data, a 100×1 Matrix{Float64}:
  2.025014193157797
  0.9414498482670517
[...]

julia> get(chn; section=:internals)[:a]
ERROR: type NamedTuple has no field a
[...]

The issue with this is that different sections can't have the same key name or else things will behave weirdly, see e.g. TuringLang/MCMCChains.jl#469. IMO this design of MCMCChains should really be improved. If we can have the same key name in different sections, then it would help to resolve this ambiguity: the parameter values could live in the :parameters section, and the pointwise logdensities could live in a different section say :lp.

[penelopeysm]

A number of months on I've pretty much come round to the idea of returning a Chains object, so this is now implemented in #1076.