I have a moderate-sized model with 20+ parameters that I am fitting in numpyro using SVI.
N_STEPS = 20_000
optimizer = numpyro.optim.Adam(step_size=0.1)
guide = autoguides.AutoLaplaceApproximation(reparam_model)
svi = SVI(reparam_model, guide, optimizer, loss=Trace_ELBO())
params = svi.run(random.PRNGKey(42), N_STEPS, *data)
However, when I try to extract the parameters to sample from the posterior, there is only one parameter in the data structure returned by svi.run
:
>>> params.params
params.params
{'auto_loc': DeviceArray([-2.163792 , -2.1976624 , -0.04213875, ..., 0.07308967,
0.07306584, 0.07304902], dtype=float32)}
What is going on here?
Thanks.