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