Code 4.54 Returns Unexpected Output for me

[cgarciga]

Hi there,

I was trying to replicate the following Code 4.54 and I get the following:

Where the expected result is:

Please find a complete and inclusive minimum viable example below, which produces the same unexpected output:

[cgarciga]

Complete text of minimum viable code below:

Howell1 = pd.read_csv("./data/Howell1.csv", sep=';').query("age>=18")

xbar = Howell1['weight'].mean()

# fit model
def model(weight, height):
    a = numpyro.sample("a", dist.Normal(178, 20))
    b = numpyro.sample("b", dist.LogNormal(0, 1))
    sigma = numpyro.sample("sigma", dist.Uniform(0, 50))
    mu = numpyro.deterministic("mu", a + b * (Howell1['weight'].values - xbar))
    numpyro.sample("height", dist.Normal(mu, sigma), obs=Howell1['height'].values)


m4_3 = AutoLaplaceApproximation(model)
svi = SVI(
    model, m4_3, optim.Adam(1), Trace_ELBO(), weight=Howell1['weight'].values, height=Howell1['height'].values
)
svi_result = svi.run(random.PRNGKey(0), 2000)
p4_3 = svi_result.params

post = m4_3.sample_posterior(random.PRNGKey(1), p4_3, (1000,))

# define sequence of weights to compute predictions for
# these values will be on the horizontal axis
weight_seq = jnp.arange(start=25, stop=71, step=1)

# use predictive to compute mu
# for each sample from posterior
# and for each weight in weight_seq
mu = Predictive(m4_3.model, post, return_sites=["mu"])(
    random.PRNGKey(2), weight_seq, None
)["mu"]
mu.shape

[Du Phan]

Oops, in a recent release, we also return deterministic sites in autoguide. I think you can do this to make the code work

post = ...
post.pop("mu")

[Du Phan]

Oh, I think I see the issue now. You used

Howell1['weight'].values

at the deterministic site mu. Please change it to weight to get the expected result.