Overly aggressive concretisation in `bundle_samples`

[penelopeysm]

Turing.jl/src/mcmc/Inference.jl

Lines 364 to 365 in 0a6a10b

	# Conretize the array before giving it to MCMCChains.
	parray = MCMCChains.concretize(parray)

This function MCMCChains.concretize is quite aggressive, it doesn't just concretise but also casts:

julia> x = Matrix{Real}(undef, 2, 0)
2×0 Matrix{Real}

julia> y = hcat(x, Float64[1.0, 2.0], Bool[true, false])
2×2 Matrix{Real}:
 1.0   true
 2.0  false

julia> map(identity, y)
2×2 Matrix{Real}:
 1.0   true
 2.0  false

julia> MCMCChains.concretize(y)
2×2 Matrix{Float64}:
 1.0  1.0
 2.0  0.0

In the context of Turing it means that literally everything gets converted into Float64 which IMO is not right

For example, Int- or Bool-valued parameters:

julia> @model function f()
           x ~ Normal()
           y ~ Poisson(1.0)
       end
f (generic function with 2 methods)

julia> chn = sample(f(), MH(), 10; progress=false);

julia> chn[:y] # should be Int
2-dimensional AxisArray{Float64,2,...} with axes:
    :iter, 1:1:10
    :chain, 1:1
And data, a 10×1 Matrix{Float64}:
 2.0
 0.0
 0.0
[...]

and also sampling statistics:

julia> @model g() = x ~ Normal()
g (generic function with 2 methods)

julia> chn = sample(g(), NUTS(), 10; progress=false, verbose=false);

julia> chn[:is_accept] # should be Bool
2-dimensional AxisArray{Float64,2,...} with axes:
    :iter, 6:1:15
    :chain, 1:1
And data, a 10×1 Matrix{Float64}:
 1.0
 1.0
[...]

julia> chn[:n_steps] # should be Int
2-dimensional AxisArray{Float64,2,...} with axes:
    :iter, 6:1:15
    :chain, 1:1
And data, a 10×1 Matrix{Float64}:
 1.0
 1.0
 1.0
[...]

[penelopeysm]

I think the real underlying problem is that if you don't attempt to concretise it, then the underlying AxisArray storage of MCMCChains is e.g. Vector{Real}.

julia> x = Matrix{Real}(undef, 2, 0)
2×0 Matrix{Real}

julia> y = hcat(x, Int[1, 2], Bool[true, false])
2×2 Matrix{Real}:
 1   true
 2  false

julia> chn = MCMCChains.Chains(y, [:int, :bool])
Chains MCMC chain (2×2×1 Array{Real, 3}):

Iterations        = 1:1:2
Number of chains  = 1
Samples per chain = 2
parameters        = int, bool

Use `describe(chains)` for summary statistics and quantiles.

Then any time you index into MCMCChains you will get Vector{Real}.

julia> chn[:int]
2-dimensional AxisArray{Real,2,...} with axes:
    :iter, 1:1:2
    :chain, 1:1
And data, a 2×1 Matrix{Real}:
 1
 2

Of course you can concretise it yourself after indexing

julia> map(identity, chn[:int])
2-dimensional AxisArray{Int64,2,...} with axes:
    :iter, 1:1:2
    :chain, 1:1
And data, a 2×1 Matrix{Int64}:
 1
 2

but MCMCChains doesn't do that automatically. So if you want chn[param] to return a concrete eltype you need to cast the entire AxisArray to a single type, which happens to be Float64.