Dynamic models fail at chain concatenation

[mhauru]

MWE from the Turing test suite, just run with multiple chains:

using Turing, Random
using Turing.RandomMeasures: DirichletProcess, ChineseRestaurantProcess
import ReverseDiff

Random.seed!(23)

@model function imm(y, alpha, ::Type{M}=Vector{Float64}) where {M}
    N = length(y)
    rpm = DirichletProcess(alpha)

    z = zeros(Int, N)
    cluster_counts = zeros(Int, N)
    fill!(cluster_counts, 0)

    for i in 1:N
        z[i] ~ ChineseRestaurantProcess(rpm, cluster_counts)
        cluster_counts[z[i]] += 1
    end

    Kmax = findlast(!iszero, cluster_counts)
    m = M(undef, Kmax)
    for k in 1:Kmax
        m[k] ~ Normal(1.0, 1.0)
    end
end

num_zs = 100
num_samples = 100
num_chains = 6
model = imm(Random.randn(num_zs), 1.0)
adbackend = AutoReverseDiff()
chn = sample(model, Gibbs(PG(10, :z), HMC(0.01, 4, :m; adtype=adbackend)), MCMCThreads(), num_samples, num_chains)

This fails with

Sampling (6 threads) 100%|█████████████████████████████████████████████████████| Time: 0:00:08
ERROR: ArgumentError: chain names differ
Stacktrace:
  [1] _cat(::Val{…}, c1::Chains{…}, args::Chains{…})
    @ MCMCChains ~/.julia/packages/MCMCChains/zFCJy/src/chains.jl:809
  [2] chainscat
    @ ~/.julia/packages/MCMCChains/zFCJy/src/chains.jl:753 [inlined]
  [3] _mapreduce(f::typeof(identity), op::typeof(chainscat), ::IndexLinear, A::Vector{Chains{…}})
    @ Base ./reduce.jl:440
  [4] _mapreduce_dim
    @ ./reducedim.jl:367 [inlined]
  [5] mapreduce
    @ ./reducedim.jl:359 [inlined]
  [6] reduce
    @ ./reducedim.jl:408 [inlined]
  [7] chainsstack(c::Vector{Chains{Union{…}, AxisArrays.AxisArray{…}, Missing, @NamedTuple{…}, @NamedTuple{…}}})
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/FSyVk/src/interface.jl:20
  [8] mcmcsample(rng::TaskLocalRNG, model::DynamicPPL.Model{…}, sampler::DynamicPPL.Sampler{…}, ::MCMCThreads, N::Int64, nchains::Int64; progress::Bool, progressname::String, initial_params::Nothing, initial_state::Nothing, kwargs::@Kwargs{…})
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/FSyVk/src/sample.jl:481
  [9] sample(rng::TaskLocalRNG, model::DynamicPPL.Model{…}, sampler::DynamicPPL.Sampler{…}, ensemble::MCMCThreads, N::Int64, n_chains::Int64; chain_type::Type, progress::Bool, kwargs::@Kwargs{})
    @ Turing.Inference ~/.julia/packages/Turing/Z4MFH/src/mcmc/Inference.jl:348
 [10] sample
    @ ~/.julia/packages/Turing/Z4MFH/src/mcmc/Inference.jl:337 [inlined]
 [11] #sample#6
    @ ~/.julia/packages/Turing/Z4MFH/src/mcmc/Inference.jl:332 [inlined]
 [12] sample
    @ ~/.julia/packages/Turing/Z4MFH/src/mcmc/Inference.jl:321 [inlined]
 [13] #sample#5
    @ ~/.julia/packages/Turing/Z4MFH/src/mcmc/Inference.jl:316 [inlined]
 [14] sample(model::DynamicPPL.Model{…}, alg::Gibbs{…}, ensemble::MCMCThreads, N::Int64, n_chains::Int64)
    @ Turing.Inference ~/.julia/packages/Turing/Z4MFH/src/mcmc/Inference.jl:308
 [15] top-level scope
    @ REPL[37]:1
Some type information was truncated. Use `show(err)` to see complete types.

At the very least this should fail more informatively.

[penelopeysm]

Slimmer repro:

using Random
using Turing
using LinearAlgebra: I

@model function f(p)
    x ~ Categorical([1 - p, p])
    y ~ MvNormal(zeros(x), I)
    return (x=x, y=y)
end

Random.seed!(468)
chn = sample(f(0.04), Prior(), MCMCSerial(), 10, 2)

The problem is that in one of the chains, the value x = 2 is never realised and so parameter names are x, y[1]; but in the other chain, it does happen to get x = 2, and the parameter names are x, y[1], y[2].

[penelopeysm]

FWIW

julia> chn = sample(Xoshiro(468), f(0.04), Prior(), MCMCSerial(), 10, 2)
ERROR: ArgumentError: chain names differ
[...]

julia> chn = sample(Xoshiro(468), f(0.04), Prior(), MCMCSerial(), 10, 2; chain_type=FlexiChains.VNChain)
FlexiChain | 10 iterations (1:10) | 2 chains (1:2)
Parameter type   VarName
Parameters       x, y
Extra keys       :logprior, :loglikelihood, :logjoint

julia> chn[@varname(y)]
┌ 10×2 DimArray{Vector{Float64}, 2} ┐
├───────────────────────────────────┴────────────────── dims ┐
  ↓ iter Sampled{Int64} 1:10 ForwardOrdered Regular Points,
  → chain Sampled{Int64} 1:2 ForwardOrdered Regular Points
└────────────────────────────────────────────────────────────┘
  ↓ →  1             2
  1     [0.229583]    [1.4895]
  2     [-0.160152]   [1.35816]
  3     [1.01454]     [0.38654]
  4     [0.255382]    [1.96437]
  5     [-0.377391]   [-0.71801]
  6     [0.475966]    [0.139307]
  7     [-0.278329]   [-1.42268]
  8     [0.0922477]   [-0.239177, 0.392541]
  9     [1.57054]     [-0.309421]
 10     [1.98705]     [-1.71261]