I’m interested in using numpyro to estimate models with “nested” inferences, similar to the nesting structure in the RSA examples with Pyro. I’ve run into some trouble trying to implement a simplified version of this. The error I eventually get is, AssertionError: all sites must have unique names but got mu duplicated.
So I’m wondering, how do I manage the naming of sample sites when inference is nested inside inference?
Minimal example
For a sort of minimal example, say we have a “meter” that takes inputs x and gives an estimate of the population mean based on those inputs. The meter might be biased, which is to say it might have a prior about what mean values are likely. We want to do BDA to infer the prior parameters of the meter.
So if we give it x = [-1,0,1,.5] and it outputs mu_est = -.50, then it has some negative bias (an unbiased estimate is .125).
# generate a vector of observations to be fed into the meter
x = dist.Normal(jnp.ones(10),1).sample(PRNGKey(123))
# define the meter and allow it to be passed a prior parameter
def meterModel(prior_mu, x=None):
mu = numpyro.sample("mu", dist.Normal(prior_mu, .5))
sigma = 1.
with numpyro.plate("obs", x.shape[0]):
numpyro.sample("obs_x", dist.Normal(mu, sigma), obs=x)
# do MCMC with this model with mu_prior = -1 and see it produces a biased estimate
mu_prior = -1.
kernel = NUTS(meterModel, target_accept_prob=.80)
posterior = MCMC(kernel, 1_000, 1_000, num_chains=1)
posterior.run(PRNGKey(0), mu_prior, x)
posterior.print_summary()
Now I want to do nested inference to infer the prior that this meter must be using to generate the biased outputs that it does.
The approach I’ve tried is to use laplace approximation for the nested inference model, and use MCMC on the outer model.
def laplace_approx(model, *args):
guide = AutoLaplaceApproximation(model)
optimizer = numpyro.optim.Minimize()
svi = SVI(model, guide, optimizer, Trace_ELBO())
init_state = svi.init(PRNGKey(1), *args)
optimal_state, loss = svi.update(init_state, *args)
return guide.get_posterior(svi.get_params(optimal_state))
def outerModel(x, meter_output=None):
prior_mu = numpyro.sample("prior_mu", dist.Normal(0,3))
meter_dist = laplace_approx(meterModel, prior_mu, x)
numpyro.sample("meter_samp", meter_dist, obs = meter_output)
kernel = NUTS(outerModel, target_accept_prob=.80)
posterior = MCMC(kernel, 250, 250, num_chains=1)
posterior.run(PRNGKey(0), x, jnp.array(.22))
But when I try this I get the error:
AssertionError: all sites must have unique names but got `mu` duplicated
I wouldn’t be surprised if there was more wrong with this code, but I am stuck. How do I separate/manage the naming of the samples sites between the inner and outer model?
