def test_relaxed_overdispersed_binomial(overdispersion): total_count = torch.arange(1, 33) probs = torch.linspace(0.1, 0.9, 16).unsqueeze(-1) d1 = binomial_dist(total_count, probs, overdispersion=overdispersion) assert isinstance(d1, dist.ExtendedBetaBinomial) with set_relaxed_distributions(): d2 = binomial_dist(total_count, probs, overdispersion=overdispersion) assert isinstance(d2, dist.Normal) assert_close(d2.mean, d1.mean) assert_close(d2.variance, d1.variance.clamp(min=_RELAX_MIN_VARIANCE))
def test_relaxed_beta_binomial(): total_count = torch.arange(1, 17) concentration1 = torch.logspace(-1, 2, 8).unsqueeze(-1) concentration0 = concentration1.unsqueeze(-1) d1 = beta_binomial_dist(concentration1, concentration0, total_count) assert isinstance(d1, dist.ExtendedBetaBinomial) with set_relaxed_distributions(): d2 = beta_binomial_dist(concentration1, concentration0, total_count) assert isinstance(d2, dist.Normal) assert_close(d2.mean, d1.mean) assert_close(d2.variance, d1.variance.clamp(min=_RELAX_MIN_VARIANCE))