Pyro changes dimension of Discrete latent variable when using NUTS (MCMC) sampler

Thanks for taking time to read my issue as given below.

The issue I need help with is that the dimensions of my Binomial distribution output changes (automatically) during the second iteration when I run the model using NUTS sampler. Because of this the remaining of my code (not given here) throws a dimension mismatch error. If I run the model function only by just calling the function (without using Sampler) it works great, even if I keep calling the function repeatedly. But it fails when I use Sampler.

I replicated the issue using a smaller and simpler code as mentioned below (this code doesn't represent my actual code but replicates the issue).

The packages I imported:

import pyro
import pyro.distributions as dist
import torch
import pyro.poutine as poutine
from pyro.infer import MCMC, NUTS

The version of Pyro is 1.5 and PyTorch is 1.7

The Model

def model ():
        
    print("***** Start ****")
    prior = torch.ones(5) / 5
    print("Prior", prior)
    
    a = pyro.sample("a", dist.Binomial(1, prior))
    print("A", a)
    
    b = pyro.sample("b", dist.Binomial(1, a)) 
    print("B", b)
    
    print("***** End *****")
    
    return b

def conditioned_model(model, data):
    print("**** Condition Model **** ")
    return poutine.condition(model, data = {"b":data})()

data = model()

Output when I call the model directly to generate simulated data

***** Start ****
Prior tensor([0.2000, 0.2000, 0.2000, 0.2000, 0.2000])
A tensor([0., 1., 0., 0., 0.])
B tensor([0., 1., 0., 0., 0.])
***** End *****

MCMC Sampler code

nuts_kernel = NUTS(conditioned_model, jit_compile=False)
mcmc = MCMC(nuts_kernel,
            num_samples=1,
            warmup_steps=1,
            num_chains=1)
mcmc.run(model, data)

Output when I run MCMC sampler (above code)

Warmup:   0%|          | 0/2 [00:00, ?it/s]
**** Condition Model **** 
***** Start ****
Prior tensor([0.2000, 0.2000, 0.2000, 0.2000, 0.2000])
A tensor([1., 0., 0., 0., 1.])
B tensor([0., 1., 0., 0., 0.])
***** End *****
**** Condition Model **** 
***** Start ****
Prior tensor([0.2000, 0.2000, 0.2000, 0.2000, 0.2000])
A tensor([0., 1.])
B tensor([0., 1., 0., 0., 0.])
***** End *****

In the above output please observe the dimension of variable A. Initially it has size 5 and later it becomes 2. Due to this remaining of my code in DINA model gives error.

In above code, variable A is based on the prior variable and the dimension of prior is 5. Then as I understand, A should always be 5. Please help me understand why it changes to 2 and how I can avoid this from happening.

Also, what I am not able to understand is that the dimension of B always remains 5. In above code, B is taking A as input, but B doesn't change the dimension even when when A changes its dimension.

Thanks a lot for the help.

Solution

I found another discussion regarding this issue.

It seems to me that the issue in my code is that NUTS try to integrate out Discrete random variables. Hence, I cannot apply a conditional flow based on the discrete random variable. See here for more information: Error with NUTS when random variable is used in control flow