def test_dirichlet_categorical(jit):
def model(data):
concentration = torch.tensor([1.0, 1.0, 1.0])
p_latent = pyro.sample('p_latent', dist.Dirichlet(concentration))
pyro.sample("obs", dist.Categorical(p_latent), obs=data)
return p_latent
true_probs = torch.tensor([0.1, 0.6, 0.3])
data = dist.Categorical(true_probs).sample(sample_shape=(torch.Size((2000,))))
hmc_kernel = HMC(model, trajectory_length=1, jit_compile=jit, ignore_jit_warnings=True)
mcmc_run = MCMC(hmc_kernel, num_samples=200, warmup_steps=100).run(data)
posterior = mcmc_run.marginal('p_latent').empirical['p_latent']
assert_equal(posterior.mean, true_probs, prec=0.02)
def test_gamma_normal(jit):
def model(data):
rate = torch.tensor([1.0, 1.0])
concentration = torch.tensor([1.0, 1.0])
p_latent = pyro.sample('p_latent', dist.Gamma(rate, concentration))
pyro.sample("obs", dist.Normal(3, p_latent), obs=data)
return p_latent
true_std = torch.tensor([0.5, 2])
data = dist.Normal(3, true_std).sample(sample_shape=(torch.Size((2000,))))
hmc_kernel = HMC(model, trajectory_length=1, step_size=0.03, adapt_step_size=False,
jit_compile=jit, ignore_jit_warnings=True)
mcmc_run = MCMC(hmc_kernel, num_samples=200, warmup_steps=200).run(data)
posterior = mcmc_run.marginal(['p_latent']).empirical['p_latent']
assert_equal(posterior.mean, true_std, prec=0.05)
def test_beta_bernoulli(jit):
def model(data):
alpha = torch.tensor([1.1, 1.1])
beta = torch.tensor([1.1, 1.1])
p_latent = pyro.sample('p_latent', dist.Beta(alpha, beta))
with pyro.plate("data", data.shape[0], dim=-2):
pyro.sample('obs', dist.Bernoulli(p_latent), obs=data)
return p_latent
true_probs = torch.tensor([0.9, 0.1])
data = dist.Bernoulli(true_probs).sample(sample_shape=(torch.Size((1000,))))
hmc_kernel = HMC(model, trajectory_length=1, max_plate_nesting=2,
jit_compile=jit, ignore_jit_warnings=True)
mcmc_run = MCMC(hmc_kernel, num_samples=800, warmup_steps=500).run(data)
posterior = mcmc_run.marginal(["p_latent"]).empirical["p_latent"]
assert_equal(posterior.mean, true_probs, prec=0.05)
def test_beta_bernoulli(step_size, adapt_step_size, adapt_mass_matrix,
full_mass):
def model(data):
alpha = torch.tensor([1.1, 1.1])
beta = torch.tensor([1.1, 1.1])
p_latent = pyro.sample("p_latent", dist.Beta(alpha, beta))
pyro.sample("obs", dist.Bernoulli(p_latent), obs=data)
return p_latent
true_probs = torch.tensor([0.9, 0.1])
data = dist.Bernoulli(true_probs).sample(
sample_shape=(torch.Size((1000, ))))
nuts_kernel = NUTS(model, step_size, adapt_step_size, adapt_mass_matrix,
full_mass)
mcmc_run = MCMC(nuts_kernel, num_samples=500, warmup_steps=100).run(data)
posterior = mcmc_run.marginal(sites='p_latent').empirical['p_latent']
assert_equal(posterior.mean, true_probs, prec=0.02)
def test_logistic_regression(step_size, trajectory_length, num_steps,
adapt_step_size, adapt_mass_matrix, full_mass):
dim = 3
data = torch.randn(2000, dim)
true_coefs = torch.arange(1., dim + 1.)
labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample()
def model(data):
coefs_mean = torch.zeros(dim)
coefs = pyro.sample('beta', dist.Normal(coefs_mean, torch.ones(dim)))
y = pyro.sample('y', dist.Bernoulli(logits=(coefs * data).sum(-1)), obs=labels)
return y
hmc_kernel = HMC(model, step_size, trajectory_length, num_steps,
adapt_step_size, adapt_mass_matrix, full_mass)
mcmc_run = MCMC(hmc_kernel, num_samples=500, warmup_steps=100, disable_progbar=True).run(data)
beta_posterior = mcmc_run.marginal(['beta']).empirical['beta']
assert_equal(rmse(true_coefs, beta_posterior.mean).item(), 0.0, prec=0.1)
def test_bernoulli_latent_model(jit):
def model(data):
y_prob = pyro.sample("y_prob", dist.Beta(1.0, 1.0))
y = pyro.sample("y", dist.Bernoulli(y_prob))
with pyro.plate("data", data.shape[0]):
z = pyro.sample("z", dist.Bernoulli(0.65 * y + 0.1))
pyro.sample("obs", dist.Normal(2. * z, 1.), obs=data)
pyro.sample("nuisance", dist.Bernoulli(0.3))
N = 2000
y_prob = torch.tensor(0.3)
y = dist.Bernoulli(y_prob).sample(torch.Size((N,)))
z = dist.Bernoulli(0.65 * y + 0.1).sample()
data = dist.Normal(2. * z, 1.0).sample()
hmc_kernel = HMC(model, trajectory_length=1, max_plate_nesting=1,
jit_compile=jit, ignore_jit_warnings=True)
mcmc_run = MCMC(hmc_kernel, num_samples=600, warmup_steps=200).run(data)
posterior = mcmc_run.marginal("y_prob").empirical["y_prob"].mean
assert_equal(posterior, y_prob, prec=0.05)
def test_logistic_regression(jit, use_multinomial_sampling):
dim = 3
data = torch.randn(2000, dim)
true_coefs = torch.arange(1., dim + 1.)
labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample()
def model(data):
coefs_mean = torch.zeros(dim)
coefs = pyro.sample('beta', dist.Normal(coefs_mean, torch.ones(dim)))
y = pyro.sample('y',
dist.Bernoulli(logits=(coefs * data).sum(-1)),
obs=labels)
return y
nuts_kernel = NUTS(model,
use_multinomial_sampling=use_multinomial_sampling,
jit_compile=jit,
ignore_jit_warnings=True)
mcmc_run = MCMC(nuts_kernel, num_samples=500, warmup_steps=100).run(data)
posterior = mcmc_run.marginal('beta').empirical['beta']
assert_equal(rmse(true_coefs, posterior.mean).item(), 0.0, prec=0.1)