def test_rejector(rate, factor):
num_samples = 100000
rates = torch.tensor(torch.tensor(rate).expand(num_samples, 1), requires_grad=True)
factors = torch.tensor(torch.tensor(factor).expand(num_samples, 1), requires_grad=True)
dist1 = Exponential(rates)
dist2 = RejectionExponential(rates, factors) # implemented using Rejector
x1 = dist1.rsample()
x2 = dist2.rsample()
assert_equal(x1.mean(), x2.mean(), prec=0.02, msg='bug in .rsample()')
assert_equal(x1.std(), x2.std(), prec=0.02, msg='bug in .rsample()')
assert_equal(dist1.log_prob(x1), dist2.log_prob(x1), msg='bug in .log_prob()')
def test_exponential_elbo(rate, factor):
num_samples = 100000
rates = torch.full((num_samples, 1), rate).requires_grad_()
factors = torch.full((num_samples, 1), factor).requires_grad_()
model = Exponential(torch.ones(num_samples, 1))
guide1 = Exponential(rates)
guide2 = RejectionExponential(rates, factors) # implemented using Rejector
grads = []
for guide in [guide1, guide2]:
grads.append(compute_elbo_grad(model, guide, [rates])[0])
expected, actual = grads
assert_equal(actual.mean(), expected.mean(), prec=0.05, msg='bad grad for rate')
actual = compute_elbo_grad(model, guide2, [factors])[0]
assert_equal(actual.mean().item(), 0.0, prec=0.05, msg='bad grad for factor')
def test_rejection_exponential_sample_shape(sample_shape, batch_shape):
rates = torch.ones(batch_shape)
factors = torch.ones(batch_shape) * 0.5
dist = RejectionExponential(rates, factors)
x = dist.rsample(sample_shape)
assert x.shape == sample_shape + batch_shape
def test_rejection_exponential_sample_shape(sample_shape, batch_shape):
rates = torch.ones(batch_shape)
factors = torch.ones(batch_shape) * 0.5
dist = RejectionExponential(rates, factors)
x = dist.rsample(sample_shape)
assert x.shape == sample_shape + batch_shape