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