numpyro.infer.Predictive returning identical samples

The following minimal example runs SVI on a simple linear model to fit 20 data points. It then samples 100 predictions (for each of the 20 inputs) from the fitted model using the guide’s posterior distribution.

However, all 100 samples (for each of the 20 inputs) are identical; that is, the empirical posterior distribution over y has no variance. I’m wondering if this is supposed to be the case, and if so, how could I modify the code to properly sample from the posterior distribution of parameters?

import numpy as np
import numpyro
import numpyro.distributions as dist
import numpyro.optim as optim
from numpyro.infer import SVI, Trace_ELBO, Predictive
from numpyro.infer.autoguide import AutoNormal
from jax import random

x = np.expand_dims(np.linspace(start=-1, stop=1, num=20), 1)
y = 3 * x + 1 + np.random.normal(size=(20, 1))

def model(x, y):
    a = numpyro.sample("a", dist.Normal(1, 1))
    b = numpyro.sample("b", dist.Normal(0, 1))
    sigma = numpyro.sample("sigma", dist.Uniform(0., 5.))
    mu = a * x + b
    numpyro.sample("y", dist.Normal(mu, sigma), obs=y)

guide = AutoNormal(model)
svi = SVI(model, guide, optim.Adam(1), Trace_ELBO())
svi_result = svi.run(random.PRNGKey(0), 1000, x, y)

params = svi_result.params
samples = guide.sample_posterior(random.PRNGKey(1), params, (100,))
predictive = Predictive(model, samples)
samples_predictive = predictive(random.PRNGKey(0), x, y)
print(samples_predictive['y'].T)

I think you need predictive(random.PRNGKey(0), x, None). Otherwise, there is no need to sample y - the output is just a collection of the input data y.

1 Like

That makes perfect sense and solved my problem–Thanks for the quick response!