def test_beta_binomial_dependent_sample():
total = 10
counts = dist.Binomial(total, 0.3).sample()
concentration1 = torch.tensor(0.5)
concentration0 = torch.tensor(1.5)
prior = dist.Beta(concentration1, concentration0)
posterior = dist.Beta(concentration1 + counts,
concentration0 + total - counts)
def model(counts):
prob = pyro.sample("prob", prior)
pyro.sample("counts", dist.Binomial(total, prob), obs=counts)
reparam_model = poutine.reparam(
model,
{
"prob":
ConjugateReparam(
lambda counts: dist.Beta(1 + counts, 1 + total - counts)),
},
)
with poutine.trace() as tr, pyro.plate("particles", 10000):
reparam_model(counts)
samples = tr.trace.nodes["prob"]["value"]
assert_close(samples.mean(), posterior.mean, atol=0.01)
assert_close(samples.std(), posterior.variance.sqrt(), atol=0.01)
def test_beta_binomial_hmc():
num_samples = 1000
total = 10
counts = dist.Binomial(total, 0.3).sample()
concentration1 = torch.tensor(0.5)
concentration0 = torch.tensor(1.5)
prior = dist.Beta(concentration1, concentration0)
likelihood = dist.Beta(1 + counts, 1 + total - counts)
posterior = dist.Beta(concentration1 + counts,
concentration0 + total - counts)
def model():
prob = pyro.sample("prob", prior)
pyro.sample("counts", dist.Binomial(total, prob), obs=counts)
reparam_model = poutine.reparam(model,
{"prob": ConjugateReparam(likelihood)})
kernel = HMC(reparam_model)
samples = MCMC(kernel, num_samples, warmup_steps=0).run()
pred = Predictive(reparam_model, samples, num_samples=num_samples)
trace = pred.get_vectorized_trace()
samples = trace.nodes["prob"]["value"]
assert_close(samples.mean(), posterior.mean, atol=0.01)
assert_close(samples.std(), posterior.variance.sqrt(), atol=0.01)
def test_init():
total = 10
counts = dist.Binomial(total, 0.3).sample()
concentration1 = torch.tensor(0.5)
concentration0 = torch.tensor(1.5)
prior = dist.Beta(concentration1, concentration0)
likelihood = dist.Beta(1 + counts, 1 + total - counts)
def model():
x = pyro.sample("x", prior)
pyro.sample("counts", dist.Binomial(total, x), obs=counts)
return x
check_init_reparam(model, ConjugateReparam(likelihood))
def test_stable_hmm_smoke(batch_shape, num_steps, hidden_dim, obs_dim):
init_dist = random_stable(batch_shape + (hidden_dim, )).to_event(1)
trans_mat = torch.randn(batch_shape + (num_steps, hidden_dim, hidden_dim),
requires_grad=True)
trans_dist = random_stable(batch_shape +
(num_steps, hidden_dim)).to_event(1)
obs_mat = torch.randn(batch_shape + (num_steps, hidden_dim, obs_dim),
requires_grad=True)
obs_dist = random_stable(batch_shape + (num_steps, obs_dim)).to_event(1)
data = obs_dist.sample()
assert data.shape == batch_shape + (num_steps, obs_dim)
def model(data):
hmm = dist.LinearHMM(init_dist,
trans_mat,
trans_dist,
obs_mat,
obs_dist,
duration=num_steps)
with pyro.plate_stack("plates", batch_shape):
z = pyro.sample("z", hmm)
pyro.sample("x", dist.Normal(z, 1).to_event(2), obs=data)
# Test that we can combine these two reparameterizers.
reparam_model = poutine.reparam(
model,
{
"z":
LinearHMMReparam(StableReparam(), StableReparam(),
StableReparam()),
},
)
reparam_model = poutine.reparam(
reparam_model,
{
"z": ConjugateReparam(dist.Normal(data, 1).to_event(2)),
},
)
reparam_guide = AutoDiagonalNormal(
reparam_model) # Models auxiliary variables.
# Smoke test only.
elbo = Trace_ELBO(num_particles=5, vectorize_particles=True)
loss = elbo.differentiable_loss(reparam_model, reparam_guide, data)
params = [trans_mat, obs_mat]
torch.autograd.grad(loss, params, retain_graph=True)
def test_gaussian_hmm_elbo(batch_shape, num_steps, hidden_dim, obs_dim):
init_dist = random_mvn(batch_shape, hidden_dim)
trans_mat = torch.randn(batch_shape + (num_steps, hidden_dim, hidden_dim),
requires_grad=True)
trans_dist = random_mvn(batch_shape + (num_steps, ), hidden_dim)
obs_mat = torch.randn(batch_shape + (num_steps, hidden_dim, obs_dim),
requires_grad=True)
obs_dist = random_mvn(batch_shape + (num_steps, ), obs_dim)
data = obs_dist.sample()
assert data.shape == batch_shape + (num_steps, obs_dim)
prior = dist.GaussianHMM(init_dist, trans_mat, trans_dist, obs_mat,
obs_dist)
likelihood = dist.Normal(data, 1).to_event(2)
posterior, log_normalizer = prior.conjugate_update(likelihood)
def model(data):
with pyro.plate_stack("plates", batch_shape):
z = pyro.sample("z", prior)
pyro.sample("x", dist.Normal(z, 1).to_event(2), obs=data)
def guide(data):
with pyro.plate_stack("plates", batch_shape):
pyro.sample("z", posterior)
reparam_model = poutine.reparam(model, {"z": ConjugateReparam(likelihood)})
def reparam_guide(data):
pass
elbo = Trace_ELBO(num_particles=1000, vectorize_particles=True)
expected_loss = elbo.differentiable_loss(model, guide, data)
actual_loss = elbo.differentiable_loss(reparam_model, reparam_guide, data)
assert_close(actual_loss, expected_loss, atol=0.01)
params = [trans_mat, obs_mat]
expected_grads = torch.autograd.grad(expected_loss,
params,
retain_graph=True)
actual_grads = torch.autograd.grad(actual_loss, params, retain_graph=True)
for a, e in zip(actual_grads, expected_grads):
assert_close(a, e, rtol=0.01)
def test_beta_binomial_elbo():
total = 10
counts = dist.Binomial(total, 0.3).sample()
concentration1 = torch.tensor(0.5, requires_grad=True)
concentration0 = torch.tensor(1.5, requires_grad=True)
prior = dist.Beta(concentration1, concentration0)
likelihood = dist.Beta(1 + counts, 1 + total - counts)
posterior = dist.Beta(concentration1 + counts,
concentration0 + total - counts)
def model():
prob = pyro.sample("prob", prior)
pyro.sample("counts", dist.Binomial(total, prob), obs=counts)
def guide():
pyro.sample("prob", posterior)
reparam_model = poutine.reparam(model,
{"prob": ConjugateReparam(likelihood)})
def reparam_guide():
pass
elbo = Trace_ELBO(num_particles=10000,
vectorize_particles=True,
max_plate_nesting=0)
expected_loss = elbo.differentiable_loss(model, guide)
actual_loss = elbo.differentiable_loss(reparam_model, reparam_guide)
assert_close(actual_loss, expected_loss, atol=0.01)
params = [concentration1, concentration0]
expected_grads = torch.autograd.grad(expected_loss,
params,
retain_graph=True)
actual_grads = torch.autograd.grad(actual_loss, params, retain_graph=True)
for a, e in zip(actual_grads, expected_grads):
assert_close(a, e, rtol=0.01)