def test_predictive_with_guide():
data = np.array([1] * 8 + [0] * 2)
def model(data):
f = numpyro.sample("beta", dist.Beta(1., 1.))
with numpyro.plate("plate", 10):
numpyro.sample("obs", dist.Bernoulli(f), obs=data)
def guide(data):
alpha_q = numpyro.param("alpha_q", 1.0,
constraint=constraints.positive)
beta_q = numpyro.param("beta_q", 1.0,
constraint=constraints.positive)
numpyro.sample("beta", dist.Beta(alpha_q, beta_q))
svi = SVI(model, guide, optim.Adam(0.1), ELBO())
svi_state = svi.init(random.PRNGKey(1), data)
def body_fn(i, val):
svi_state, _ = svi.update(val, data)
return svi_state
svi_state = lax.fori_loop(0, 1000, body_fn, svi_state)
params = svi.get_params(svi_state)
predictive = Predictive(model, guide=guide, params=params, num_samples=1000)
obs_pred = predictive.get_samples(random.PRNGKey(2), data=None)["obs"]
assert_allclose(np.mean(obs_pred), 0.8, atol=0.05)
def test_beta_bernoulli(auto_class):
data = jnp.array([[1.0] * 8 + [0.0] * 2, [1.0] * 4 + [0.0] * 6]).T
def model(data):
f = numpyro.sample("beta", dist.Beta(jnp.ones(2), jnp.ones(2)))
numpyro.sample("obs", dist.Bernoulli(f), obs=data)
adam = optim.Adam(0.01)
guide = auto_class(model, init_loc_fn=init_strategy)
svi = SVI(model, guide, adam, Trace_ELBO())
svi_state = svi.init(random.PRNGKey(1), data)
def body_fn(i, val):
svi_state, loss = svi.update(val, data)
return svi_state
svi_state = fori_loop(0, 3000, body_fn, svi_state)
params = svi.get_params(svi_state)
true_coefs = (jnp.sum(data, axis=0) + 1) / (data.shape[0] + 2)
# test .sample_posterior method
posterior_samples = guide.sample_posterior(
random.PRNGKey(1), params, sample_shape=(1000,)
)
assert_allclose(jnp.mean(posterior_samples["beta"], 0), true_coefs, atol=0.05)
# Predictive can be instantiated from posterior samples...
predictive = Predictive(model, posterior_samples=posterior_samples)
predictive_samples = predictive(random.PRNGKey(1), None)
assert predictive_samples["obs"].shape == (1000, 2)
# ... or from the guide + params
predictive = Predictive(model, guide=guide, params=params, num_samples=1000)
predictive_samples = predictive(random.PRNGKey(1), None)
assert predictive_samples["obs"].shape == (1000, 2)
def test_prior_predictive():
model, data, true_probs = beta_bernoulli()
predictive_samples = Predictive(model, num_samples=100).get_samples(random.PRNGKey(1))
assert predictive_samples.keys() == {"beta", "obs"}
# check shapes
assert predictive_samples["beta"].shape == (100,) + true_probs.shape
assert predictive_samples["obs"].shape == (100,) + data.shape
def test_beta_bernoulli(auto_class):
data = jnp.array([[1.0] * 8 + [0.0] * 2, [1.0] * 4 + [0.0] * 6]).T
N = len(data)
def model(data):
f = numpyro.sample("beta",
dist.Beta(jnp.ones(2), jnp.ones(2)).to_event())
with numpyro.plate("N", N):
numpyro.sample("obs", dist.Bernoulli(f).to_event(1), obs=data)
adam = optim.Adam(0.01)
if auto_class == AutoDAIS:
guide = auto_class(model,
init_loc_fn=init_strategy,
base_dist="cholesky")
else:
guide = auto_class(model, init_loc_fn=init_strategy)
svi = SVI(model, guide, adam, Trace_ELBO())
svi_state = svi.init(random.PRNGKey(1), data)
def body_fn(i, val):
svi_state, loss = svi.update(val, data)
return svi_state
svi_state = fori_loop(0, 3000, body_fn, svi_state)
params = svi.get_params(svi_state)
true_coefs = (jnp.sum(data, axis=0) + 1) / (data.shape[0] + 2)
# test .sample_posterior method
posterior_samples = guide.sample_posterior(random.PRNGKey(1),
params,
sample_shape=(1000, ))
posterior_mean = jnp.mean(posterior_samples["beta"], 0)
assert_allclose(posterior_mean, true_coefs, atol=0.05)
if auto_class not in [AutoDAIS, AutoDelta, AutoIAFNormal, AutoBNAFNormal]:
quantiles = guide.quantiles(params, [0.2, 0.5, 0.8])
assert quantiles["beta"].shape == (3, 2)
# Predictive can be instantiated from posterior samples...
predictive = Predictive(model, posterior_samples=posterior_samples)
predictive_samples = predictive(random.PRNGKey(1), None)
assert predictive_samples["obs"].shape == (1000, N, 2)
# ... or from the guide + params
predictive = Predictive(model,
guide=guide,
params=params,
num_samples=1000)
predictive_samples = predictive(random.PRNGKey(1), None)
assert predictive_samples["obs"].shape == (1000, N, 2)
def test_log_likelihood(batch_shape):
model, data, _ = beta_bernoulli()
samples = Predictive(model, return_sites=["beta"], num_samples=200)(random.PRNGKey(1))
batch_size = int(np.prod(batch_shape))
samples = {'beta': samples['beta'][:batch_size].reshape(batch_shape + (1, -1))}
preds = Predictive(model, samples, batch_ndims=len(batch_shape))(random.PRNGKey(2))
loglik = log_likelihood(model, samples, data, batch_ndims=len(batch_shape))
assert preds.keys() == {"beta_sq", "obs"}
assert loglik.keys() == {"obs"}
# check shapes
assert preds["obs"].shape == batch_shape + data.shape
assert loglik["obs"].shape == batch_shape + data.shape
assert_allclose(loglik["obs"], dist.Bernoulli(samples["beta"]).log_prob(data))
def test_autoguide_deterministic(auto_class):
def model(y=None):
n = y.size if y is not None else 1
mu = numpyro.sample("mu", dist.Normal(0, 5))
sigma = numpyro.param("sigma", 1, constraint=constraints.positive)
y = numpyro.sample("y", dist.Normal(mu, sigma).expand((n,)), obs=y)
numpyro.deterministic("z", (y - mu) / sigma)
mu, sigma = 2, 3
y = mu + sigma * random.normal(random.PRNGKey(0), shape=(300,))
y_train = y[:200]
y_test = y[200:]
guide = auto_class(model)
optimiser = numpyro.optim.Adam(step_size=0.01)
svi = SVI(model, guide, optimiser, Trace_ELBO())
params, losses = svi.run(random.PRNGKey(0), num_steps=500, y=y_train)
posterior_samples = guide.sample_posterior(
random.PRNGKey(0), params, sample_shape=(1000,)
)
predictive = Predictive(model, posterior_samples, params=params)
predictive_samples = predictive(random.PRNGKey(0), y_test)
assert predictive_samples["y"].shape == (1000, 100)
assert predictive_samples["z"].shape == (1000, 100)
assert_allclose(
(predictive_samples["y"] - posterior_samples["mu"][..., None])
/ params["sigma"],
predictive_samples["z"],
atol=0.05,
)
def test_predictive_with_guide():
data = jnp.array([1] * 8 + [0] * 2)
def model(data):
f = numpyro.sample("beta", dist.Beta(1., 1.))
with numpyro.plate("plate", 10):
numpyro.deterministic("beta_sq", f**2)
numpyro.sample("obs", dist.Bernoulli(f), obs=data)
def guide(data):
alpha_q = numpyro.param("alpha_q",
1.0,
constraint=constraints.positive)
beta_q = numpyro.param("beta_q", 1.0, constraint=constraints.positive)
numpyro.sample("beta", dist.Beta(alpha_q, beta_q))
svi = SVI(model, guide, optim.Adam(0.1), Trace_ELBO())
svi_result = svi.run(random.PRNGKey(1), 3000, data)
params = svi_result.params
predictive = Predictive(model,
guide=guide,
params=params,
num_samples=1000)(random.PRNGKey(2), data=None)
assert predictive["beta_sq"].shape == (1000, )
obs_pred = predictive["obs"].astype(np.float32)
assert_allclose(jnp.mean(obs_pred), 0.8, atol=0.05)
def test_predictive(parallel):
model, data, true_probs = beta_bernoulli()
mcmc = MCMC(NUTS(model), num_warmup=100, num_samples=100)
mcmc.run(random.PRNGKey(0), data)
samples = mcmc.get_samples()
predictive = Predictive(model, samples, parallel=parallel)
predictive_samples = predictive.get_samples(random.PRNGKey(1))
assert predictive_samples.keys() == {"obs"}
predictive.return_sites = ["beta", "obs"]
predictive_samples = predictive.get_samples(random.PRNGKey(1))
# check shapes
assert predictive_samples["beta"].shape == (100,) + true_probs.shape
assert predictive_samples["obs"].shape == (100,) + data.shape
# check sample mean
assert_allclose(predictive_samples["obs"].reshape((-1,) + true_probs.shape).mean(0), true_probs, rtol=0.1)
def test_predictive_with_particles():
num_particles = 5
num_samples = 2
fdim = 3
num_data = 10
def model(x, y=None):
latent = numpyro.sample("latent",
dist.Normal(0.0, jnp.ones(fdim)).to_event(1))
with numpyro.plate("data", x.shape[0]):
numpyro.sample("y", dist.Normal(jnp.matmul(x, latent), 1.0), obs=y)
def guide(x, y=None):
latent_loc = numpyro.param("latent_loc",
jnp.ones(fdim),
constraint=constraints.real)
assert latent_loc.ndim == 1
numpyro.sample("latent", dist.Normal(latent_loc, 1.0).to_event(1))
params = {"latent_loc": jnp.zeros((num_particles, fdim))}
x = dist.Normal(jnp.full(3, 0.2), 1.0).sample(random.PRNGKey(0),
(num_data, ))
predictions = Predictive(
model,
guide=guide,
params=params,
num_samples=num_samples,
batch_ndims=1,
)(random.PRNGKey(0), x)
assert predictions["y"].shape == (num_samples, num_particles, num_data)
def test_log_likelihood():
model, data, _ = beta_bernoulli()
samples = Predictive(model, return_sites=["beta"], num_samples=100).get_samples(random.PRNGKey(1))
loglik = log_likelihood(model, samples, data)
assert loglik.keys() == {"obs"}
# check shapes
assert loglik["obs"].shape == (100,) + data.shape
assert_allclose(loglik["obs"], dist.Bernoulli(samples["beta"].reshape((100, 1, -1))).log_prob(data))
def test_prior_predictive(batch_ndims):
model, data, true_probs = beta_bernoulli()
predictive = Predictive(model, num_samples=100, batch_ndims=batch_ndims)
predictive_samples = predictive(random.PRNGKey(1))
assert predictive_samples.keys() == {"beta", "beta_sq", "obs"}
# check shapes
batch_shape = (1, ) * (batch_ndims - 1) + (100, )
assert predictive_samples["beta"].shape == batch_shape + true_probs.shape
assert predictive_samples["obs"].shape == batch_shape + data.shape
def test_predictive_with_improper():
true_coef = 0.9
def model(data):
alpha = numpyro.sample('alpha', dist.Uniform(0, 1))
loc = numpyro.param('loc', 0., constraint=constraints.interval(0., alpha))
numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
data = true_coef + random.normal(random.PRNGKey(0), (1000,))
kernel = NUTS(model=model)
mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000)
mcmc.run(random.PRNGKey(0), data)
samples = mcmc.get_samples()
obs_pred = Predictive(model, samples).get_samples(random.PRNGKey(1), data=None)["obs"]
assert_allclose(np.mean(obs_pred), true_coef, atol=0.05)
def test_predictive_with_improper():
true_coef = 0.9
def model(data):
alpha = numpyro.sample('alpha', dist.Uniform(0, 1))
with handlers.reparam(config={'loc': TransformReparam()}):
loc = numpyro.sample('loc', dist.TransformedDistribution(
dist.Uniform(0, 1).mask(False),
AffineTransform(0, alpha)))
numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
data = true_coef + random.normal(random.PRNGKey(0), (1000,))
kernel = NUTS(model=model)
mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000)
mcmc.run(random.PRNGKey(0), data)
samples = mcmc.get_samples()
obs_pred = Predictive(model, samples)(random.PRNGKey(1), data=None)["obs"]
assert_allclose(jnp.mean(obs_pred), true_coef, atol=0.05)
def test_get_mask_optimization():
def model():
with numpyro.handlers.seed(rng_seed=0):
x = numpyro.sample("x", dist.Normal(0, 1))
numpyro.sample("y", dist.Normal(x, 1), obs=0.)
called.add("model-always")
if numpyro.get_mask() is not False:
called.add("model-sometimes")
numpyro.factor("f", x + 1)
def guide():
with numpyro.handlers.seed(rng_seed=1):
x = numpyro.sample("x", dist.Normal(0, 1))
called.add("guide-always")
if numpyro.get_mask() is not False:
called.add("guide-sometimes")
numpyro.factor("g", 2 - x)
called = set()
trace = handlers.trace(guide).get_trace()
handlers.replay(model, trace)()
assert "model-always" in called
assert "guide-always" in called
assert "model-sometimes" in called
assert "guide-sometimes" in called
called = set()
with handlers.mask(mask=False):
trace = handlers.trace(guide).get_trace()
handlers.replay(model, trace)()
assert "model-always" in called
assert "guide-always" in called
assert "model-sometimes" not in called
assert "guide-sometimes" not in called
called = set()
Predictive(model, guide=guide, num_samples=2,
parallel=True)(random.PRNGKey(2))
assert "model-always" in called
assert "guide-always" in called
assert "model-sometimes" not in called
assert "guide-sometimes" not in called
