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.

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