How to obtain coefficients as well as observations during prediction in a bayesian logistic regression model?

Hi! For a simple bayesian logistic regression model:

def logistic_regression(data, labels=None, D):
    coefs = numpyro.sample('coefs', dist.Normal(jnp.zeros(D), jnp.ones(D)))
    intercept = numpyro.sample('intercept', dist.Normal(0., 10.))
    logits = jnp.sum(coefs * data + intercept, axis=-1)
    return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)

how does one get 'coefs' as well as 'obs' during the use of predictive()? For example, right now I’m trying to get 1000 predictions (samples) given a single test input via HMC. I get a list of predictions but I would also like to get, for each prediction, the coefs for that prediction. So something like a (1000, D) shaped array. Thanks.

You can pass a list of return_sites to Predictive, which as described in its docstring defaults to returning only sites which do not appear in posterior_samples:

predictive = Predictive(
    return_sites=["coefs", "intercept", "obs"]
1 Like