Ok, I now learned that MCMC sampling with a variable number of discrete latent variables requires enumeration. And from a smaller example I made I learned that NUTS seems to be performing parallel enumeration automatically, which does not work because the control flow depends on whether a new atom is drawn. I can force sequential enumeration, but:
- This does not give correct results for very small examples that can be run with parallel enumeration
- The chain immediately gets stuck on a specific number of atoms and never samples from any other possibility, defeating the purpose of allowing a variable number of hidden variables
(I also read that you can only use parallel enumeration in the model?)
I am now stuck at every end with these questions, does anyone have some insight? I could provide some code demonstrating the problems if that helps.