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