Params returned by SVI.run is missing most model parameters

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.

i believe this is as expected. the dimension of auto_loc should equal the dimension of your latent space. you can use the median and quantiles methods to extract more information

1 Like