def test_kl_divergence():
mask = torch.tensor([[0, 1], [1, 1]]).bool()
p = Normal(torch.randn(2, 2), torch.randn(2, 2).exp())
q = Normal(torch.randn(2, 2), torch.randn(2, 2).exp())
expected = kl_divergence(p.to_event(2), q.to_event(2))
actual = (kl_divergence(p.mask(mask).to_event(2),
q.mask(mask).to_event(2)) +
kl_divergence(p.mask(~mask).to_event(2),
q.mask(~mask).to_event(2)))
assert_equal(actual, expected)
def test_kl_divergence_type(p_mask, q_mask):
p = Normal(torch.randn(2, 2), torch.randn(2, 2).exp())
q = Normal(torch.randn(2, 2), torch.randn(2, 2).exp())
mask = (
(torch.tensor(p_mask) if isinstance(p_mask, bool) else p_mask) &
(torch.tensor(q_mask) if isinstance(q_mask, bool) else q_mask)).expand(
2, 2)
expected = kl_divergence(p, q)
expected[~mask] = 0
actual = kl_divergence(p.mask(p_mask), q.mask(q_mask))
if p_mask is False or q_mask is False:
assert isinstance(actual, float) and actual == 0.
else:
assert_equal(actual, expected)
def test_broadcast(event_shape, dist_shape, mask_shape):
mask = torch.empty(torch.Size(mask_shape)).bernoulli_(0.5).bool()
base_dist = Normal(torch.zeros(dist_shape + event_shape), 1.)
base_dist = base_dist.to_event(len(event_shape))
assert base_dist.batch_shape == dist_shape
assert base_dist.event_shape == event_shape
d = base_dist.mask(mask)
d_shape = broadcast_shape(mask.shape, base_dist.batch_shape)
assert d.batch_shape == d_shape
assert d.event_shape == event_shape
def test_mask_type(mask):
p = Normal(torch.randn(2, 2), torch.randn(2, 2).exp())
p_masked = p.mask(mask)
if isinstance(mask, bool):
mask = torch.tensor(mask)
x = p.sample()
actual = p_masked.log_prob(x)
expected = p.log_prob(x) * mask.float()
assert_equal(actual, expected)
actual = p_masked.score_parts(x)
expected = p.score_parts(x)
for a, e in zip(actual, expected):
if isinstance(e, torch.Tensor):
e = e * mask.float()
assert_equal(a, e)