Issue 1:
🐛 Bug

I think the Kl divergence is summing over the batch dimensions when it shouldn't, at least for the Gaussian case

Here's the transformed distribution function:

def _kl_transformed_transformed(p, q):
if p.transforms != q.transforms:
raise NotImplementedError
if p.event_shape != q.event_shape:
raise NotImplementedError
# extra_event_dim = len(p.event_shape) - len(p.base_dist.event_shape)
extra_event_dim = len(p.event_shape)
base_kl_divergence = kl_divergence(p.base_dist, q.base_dist) #call to indep_indep below
#this will again sum over kl_divergence for each entry in batch
return _sum_rightmost(base_kl_divergence, extra_event_dim)

Here's independent_independent KL

@register_kl(Independent, Independent)
def _kl_independent_independent(p, q):
shared_ndims = min(p.reinterpreted_batch_ndims, q.reinterpreted_batch_ndims)
p_ndims = p.reinterpreted_batch_ndims - shared_ndims
q_ndims = q.reinterpreted_batch_ndims - shared_ndims
p = Independent(p.base_dist, p_ndims) if p_ndims else p.base_dist
q = Independent(q.base_dist, q_ndims) if q_ndims else q.base_dist
kl = kl_divergence(p, q)
if shared_ndims:
#this line gets called when base_dist is Gaussian
kl = sum_rightmost(kl, shared_ndims)
return kl

To Reproduce

I think something like this will do it:

p = Gaussian(torch.zeros(3, 10), torch.ones(3, 10)).independent(1)
q = Gaussian(torch.ones(3,10) *2, torch.ones(3,10).independent(1)
kl = kl_divergence(p, q)
print(kl.shape) #should be something like [3, 1] but will instead output []

expected behavior

Maybe I misunderstand kl_divergence function, but I don't think it should be summing over batch.

PyTorch Version (e.g., 1.0): 1.2

OS (e.g., Linux): Ubuntu 16.04

How you installed PyTorch (conda, pip, source): conda/ pip

Build command you used (if compiling from source):

Python version: 3.6.9

CUDA/cuDNN version: 9.0

GPU models and configuration:

Any other relevant information:

I originally found this using the MeanFieldTrace(...) class in the Pyro probabilistic programming language with normalizing flows. The Kl divergence function is used underneath the hood, and checks that the kl_divergence dimensions are the same as the q.batch_shape which in my case they were not.

cc @vincentqb @fritzo @neerajprad @alicanb @vishwakftw

Issue 2:
This is a copy of an issue I posted for pytorch: pytorch/pytorch#29698

I think this is an issue with the _kl_independent_independent class which is summing over batches when it shouldn't...

Guidelines

I think the Kl divergence is summing over the batch dimensions when it shouldn't, at least for the Gaussian case

Here's the transformed distribution function:

def _kl_transformed_transformed(p, q):
if p.transforms != q.transforms:
raise NotImplementedError
if p.event_shape != q.event_shape:
raise NotImplementedError
# extra_event_dim = len(p.event_shape) - len(p.base_dist.event_shape)
extra_event_dim = len(p.event_shape)
base_kl_divergence = kl_divergence(p.base_dist, q.base_dist) #call to indep_indep below
#this will again sum over kl_divergence for each entry in batch
return _sum_rightmost(base_kl_divergence, extra_event_dim)

Here's independent_independent KL

@register_kl(Independent, Independent)
def _kl_independent_independent(p, q):
shared_ndims = min(p.reinterpreted_batch_ndims, q.reinterpreted_batch_ndims)
p_ndims = p.reinterpreted_batch_ndims - shared_ndims
q_ndims = q.reinterpreted_batch_ndims - shared_ndims
p = Independent(p.base_dist, p_ndims) if p_ndims else p.base_dist
q = Independent(q.base_dist, q_ndims) if q_ndims else q.base_dist
kl = kl_divergence(p, q)
if shared_ndims:
#this line gets called when base_dist is Gaussian
kl = sum_rightmost(kl, shared_ndims)
return kl

To Reproduce

I think something like this will do it:

p = Gaussian(torch.zeros(3, 10), torch.ones(3, 10)).independent(1)
q = Gaussian(torch.ones(3,10) *2, torch.ones(3,10).independent(1)
kl = kl_divergence(p, q)
print(kl.shape) #should be something like [3, 1] but will instead output []

expected behavior

Maybe I misunderstand kl_divergence function, but I don't think it should be summing over batch.

PyTorch Version (e.g., 1.0): 1.2

OS (e.g., Linux): Ubuntu 16.04

How you installed PyTorch (conda, pip, source): conda/ pip

Build command you used (if compiling from source):

Python version: 3.6.9

CUDA/cuDNN version: 9.0

GPU models and configuration:

Any other relevant information:

I originally found this using the MeanFieldTrace(...) class in the Pyro probabilistic programming language with normalizing flows. The Kl divergence function is used underneath the hood, and checks that the kl_divergence dimensions are the same as the q.batch_shape which in my case they were not.