test_vae_utils.py

from torch.distributions import MultivariateNormal
from torch.distributions import kl_divergence

mean = torch.rand(10, 100)
logstds = torch.rand(10, 100)
x = torch.rand(10, 100)

dist = MultivariateNormal(mean, scale_tril=torch.stack([torch.diag((ls).exp()) for ls in logstds]))
print('Testing LL')
print(dist.log_prob(x).mean())
print(gaussian_log_likelihood(x, mean, logstds).mean())

print('Testing KL')
standard_normal = MultivariateNormal(torch.zeros_like(mean), scale_tril=torch.stack([torch.eye(mean.size(1)) for _ in range(mean.size(0))]))
print(kl_divergence(dist, standard_normal).mean())
print(KL_divergence(mean, logstds).mean())