Bug: torch.distributions.mixture_same_distribution._pad_mixture_dimension

[CJMenart]

🐛 Describe the bug

I already tracked down what's going on with this one; there is a bug in mixture_same_distribution._pad_mixture_dimension. It uses numel() on Size objects in order to get the number of dimensions in its tensors--however, Size.numel() multiplies all the elements of the size tensor together (it's telling you how many numel are in the tensor whose size it represents!), thus returning what is clearly the wrong value.

This (which is in both 1.8, which I was on originally, and master):

def _pad_mixture_dimensions(self, x):

        dist_batch_ndims = self.batch_shape.numel()
        cat_batch_ndims = self.mixture_distribution.batch_shape.numel()
        pad_ndims = 0 if cat_batch_ndims == 1 else \
            dist_batch_ndims - cat_batch_ndims
        xs = x.shape
        x = x.reshape(xs[:-1] + torch.Size(pad_ndims * [1]) +
                      xs[-1:] + torch.Size(self._event_ndims * [1]))
        return x

Should be this:

def _pad_mixture_dimensions(self, x):
        dist_batch_ndims = len(self.batch_shape)
        cat_batch_ndims = len(self.mixture_distribution.batch_shape)
        pad_ndims = 0 if cat_batch_ndims == 1 else \
            dist_batch_ndims - cat_batch_ndims
        xs = x.shape
        x = x.reshape(xs[:-1] + torch.Size(pad_ndims * [1]) +
                      xs[-1:] + torch.Size(self._event_ndims * [1]))
        return x

I didn't see this error anywhere else, but I haven't really done a thorough search.
I don't know what branch it's right to make a pull request on (master?) so I thought I'd report it here first.

Versions

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] torch==1.8.0
[pip3] torchvision==0.9.0
[conda] blas 2.109 mkl conda-forge
[conda] blas-devel 3.9.0 9_mkl conda-forge
[conda] cudatoolkit 10.2.89 h8f6ccaa_8 conda-forge
[conda] ffmpeg 4.3 hf484d3e_0 pytorch
[conda] libblas 3.9.0 9_mkl conda-forge
[conda] libcblas 3.9.0 9_mkl conda-forge
[conda] liblapack 3.9.0 9_mkl conda-forge
[conda] liblapacke 3.9.0 9_mkl conda-forge
[conda] mkl 2021.2.0 h726a3e6_389 conda-forge
[conda] mkl-devel 2021.2.0 ha770c72_390 conda-forge
[conda] mkl-include 2021.2.0 h726a3e6_389 conda-forge
[conda] numpy 1.20.3 py38h9894fe3_1 conda-forge
[conda] pytorch 1.8.0 py3.8_cuda10.2_cudnn7.6.5_0 pytorch
[conda] torchvision 0.9.0 py38_cu102 pytorch

cc @ezyang @gchanan @zou3519 @fritzo @neerajprad @alicanb @nikitaved

[mruberry]

Would you extend this issue with a small code snippet that demonstrates the bug and shows the before/after with the proposed fix?

Yep, just open a pull request for your branch to be merged with master and we'll be able to review it.

[CJMenart]

Example given:

The below snippet, under the current version, will throw "RuntimeError: only tensors with up to 64 dims are supported", because the bug is trying to create a 1000-dimension tensor. It correctly prints the mean of the gaussian mixture (100 zeros) under the fix.

import torch
from torch.distributions import Categorical, Normal, MixtureSameFamily

def bug_example():
    # A 3d tensor of gaussian distributions
    component_distribution = Normal(torch.zeros([10,10,10]), torch.ones([10,10,10]))
    # A 3D tensor of uniform distributions over 10 classes, 
    mix_distribution = Categorical(torch.ones([10,1,10]))
    # This SHOULD be 10 mixtures of 10 10-dimensional isotropic Gaussians
    gmm = MixtureSameFamily(mix_distribution, component_distribution)
    mean = gmm.mean  # This kicks off the call to _pad_mixture_dimension
    print(mean)
    

if __name__ == '__main__':
    bug_example()