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_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