Sample size changes at each sample

numpyro
1

I am trying to code a simple model with a categorical sample (I minimized the problem). However, in each sample, the dimension of ind changes. I couldn’t figure out how/why it produces that or how I can solve this. I would appreciate any feedback on that.

def model(n):
    if n==1:
        return

    probs = jnp.ones(n-1) / (n-1)
    ind = numpyro.sample('ind'+'-'+str(n), dist.Categorical(probs=probs)
)
    print(ind)
    model(n-1)
rng_key = random.PRNGKey(0)
rng_key, rng_key_ = random.split(rng_key)

# Run NUTS.
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=1000, num_samples=2000)
rng_key = random.PRNGKey(3)
mcmc.run(rng_key, n=5)
samples = mcmc.get_samples()
2

https://pyro.ai/examples/enumeration.html

1 Like