当使用NUTSMCMC采样器时，Pyro会更改离散潜变量的维数

问题描述

感谢您抽出时间阅读下面的问题。

我需要帮助的问题是，当我使用NUTS采样器运行模型时，第二次迭代期间我的二项分布输出的尺寸会自动更改。因此，我的其余代码（此处未给出）会引发尺寸不匹配错误。如果仅通过调用函数（而不使用Sampler）来运行模型函数，则即使我不断重复调用该函数也能很好地工作。但是当我使用Sampler时失败。

我使用如下所述的更简单的代码复制了问题（此代码并不代表我的实际代码，而是复制了问题）。

我导入的软件包：

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

Pyro的版本为1.5，PyTorch的版本为1.7

模型

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",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()

当我直接调用模型以生成模拟数据时输出

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

MCMC采样器代码

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)

运行MCMC采样器时的输出（以上代码）

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

在上面的输出中，请注意变量 A 的尺寸。最初它的大小为5，后来变成2。由于我在DINA模型中的其余代码，导致出现错误。

在上面的代码中，变量 A 基于 prior 变量，而 prior 的维度为5。那么据我所知， A 应该始终为5。请帮助我理解为什么它变为2，以及如何避免这种情况发生。

此外，我无法理解的是， B 的尺寸始终保持为5。在上面的代码中， B 占用了 A 作为输入，但是即使 A 更改尺寸， B 也不会更改尺寸。

非常感谢您的帮助。

解决方法

我发现了关于此问题的另一讨论。

在我看来，我的代码中的问题是NUTS尝试集成离散离散变量。因此，我无法基于离散随机变量应用条件流。有关更多信息，请参见此处：Error with NUTS when random variable is used in control flow

probabilistic-programming pyro