Repository: pyro-ppl/pyro

Issue:
[bug] MixtureOfDiagNormals
### Problem description

I recently tried to use [`MixtureOfDiagNormals`](https://github.com/pyro-ppl/pyro/blob/c00bcc3fb701327b07e94de88754da49b7f29ebe/pyro/distributions/diag_normal_mixture.py#L15) and ran across two issues while using it together with an `AutoNormal` guide.

1. The distributions [`MixtureOfDiagNormal`](https://github.com/pyro-ppl/pyro/blob/c00bcc3fb701327b07e94de88754da49b7f29ebe/pyro/distributions/diag_normal_mixture.py#L15) and [`MixtureOfDiagNormalsSharedCovariance`](https://github.com/pyro-ppl/pyro/blob/c00bcc3fb701327b07e94de88754da49b7f29ebe/pyro/distributions/diag_normal_mixture_shared_cov.py#L15) do not implement the `support` property, which leads to problems when using the distribution for latent variables and e.g. the `AutoNormal` guide. This can be solved easily by adding `support = constraints.real_vector` at the beginning of both classes. I assume that `real_vector` is the correct constraint here since `D=1` is not supported anyways.
2. There seems to be a bug in the [`forward`](https://github.com/pyro-ppl/pyro/blob/c00bcc3fb701327b07e94de88754da49b7f29ebe/pyro/distributions/diag_normal_mixture.py#L134) method of `_MixDiagNormalSample`. Specifically, when creating the white noise
   ```python
       white = locs.new(noise_shape).normal_()
   ```
   This creates a new tensor of size `(len(noise_shape),)` instead of expected size `noise_shape`. This is actually done correctly in [`_MixDiagNormalSharedCovarianceSample`](https://github.com/pyro-ppl/pyro/blob/c00bcc3fb701327b07e94de88754da49b7f29ebe/pyro/distributions/diag_normal_mixture_shared_cov.py#L137C30-L137C30) with a call to `torch.randn` instead.
   The correct way to do it would be
   ```python
       white = locs.new(*noise_shape).normal_()
   ```
   However, since `torch.Tensor.new` does not appear in the PyTorch documentation (at least I didn't find it), I assume it is recommended to use a method like `torch.Tensor.new_empty` instead.
   ```python
       white = locs.new_empty(noise_shape).normal_()
   ```
   works as expected. Alternatively, the same approach as for `MixtureOfDiagNormalSharedCovariance` is possible.

### Package versions

PyTorch: version 2.0.1
Pyro: version 1.8.6



Discussion:

Resolving patch:
diff --git a/pyro/distributions/diag_normal_mixture.py b/pyro/distributions/diag_normal_mixture.py
--- a/pyro/distributions/diag_normal_mixture.py
+++ b/pyro/distributions/diag_normal_mixture.py
@@ -46,13 +46,15 @@ class MixtureOfDiagNormals(TorchDistribution):
         "coord_scale": constraints.positive,
         "component_logits": constraints.real,
     }
+    support = constraints.real_vector
 
     def __init__(self, locs, coord_scale, component_logits):
         self.batch_mode = locs.dim() > 2
         assert coord_scale.shape == locs.shape
-        assert (
-            self.batch_mode or locs.dim() == 2
-        ), "The locs parameter in MixtureOfDiagNormals should be K x D dimensional (or B x K x D if doing batches)"
+        assert self.batch_mode or locs.dim() == 2, (
+            "The locs parameter in MixtureOfDiagNormals should be K x D dimensional "
+            "(or ... x B x K x D if doing batches)"
+        )
         if not self.batch_mode:
             assert (
                 coord_scale.dim() == 2
@@ -65,10 +67,10 @@ def __init__(self, locs, coord_scale, component_logits):
         else:
             assert (
                 coord_scale.dim() > 2
-            ), "The coord_scale parameter in MixtureOfDiagNormals should be B x K x D dimensional"
+            ), "The coord_scale parameter in MixtureOfDiagNormals should be ... x B x K x D dimensional"
             assert (
                 component_logits.dim() > 1
-            ), "The component_logits parameter in MixtureOfDiagNormals should be B x K dimensional"
+            ), "The component_logits parameter in MixtureOfDiagNormals should be ... x B x K dimensional"
             assert component_logits.size(-1) == locs.size(-2)
             batch_shape = tuple(locs.shape[:-2])
 
@@ -133,7 +135,7 @@ class _MixDiagNormalSample(Function):
     @staticmethod
     def forward(ctx, locs, scales, component_logits, pis, which, noise_shape):
         dim = scales.size(-1)
-        white = locs.new(noise_shape).normal_()
+        white = locs.new_empty(noise_shape).normal_()
         n_unsqueezes = locs.dim() - which.dim()
         for _ in range(n_unsqueezes):
             which = which.unsqueeze(-1)
diff --git a/pyro/distributions/diag_normal_mixture_shared_cov.py b/pyro/distributions/diag_normal_mixture_shared_cov.py
--- a/pyro/distributions/diag_normal_mixture_shared_cov.py
+++ b/pyro/distributions/diag_normal_mixture_shared_cov.py
@@ -45,28 +45,31 @@ class MixtureOfDiagNormalsSharedCovariance(TorchDistribution):
         "coord_scale": constraints.positive,
         "component_logits": constraints.real,
     }
+    support = constraints.real_vector
 
     def __init__(self, locs, coord_scale, component_logits):
         self.batch_mode = locs.dim() > 2
-        assert (
-            self.batch_mode or locs.dim() == 2
-        ), "The locs parameter in MixtureOfDiagNormals should be K x D dimensional (or ... x B x K x D in batch mode)"
+        assert self.batch_mode or locs.dim() == 2, (
+            "The locs parameter in MixtureOfDiagNormalsSharedCovariance should be K x D dimensional "
+            "(or ... x B x K x D in batch mode)"
+        )
         if not self.batch_mode:
             assert (
                 coord_scale.dim() == 1
-            ), "The coord_scale parameter in MixtureOfDiagNormals should be D dimensional"
+            ), "The coord_scale parameter in MixtureOfDiagNormalsSharedCovariance should be D dimensional"
             assert (
                 component_logits.dim() == 1
-            ), "The component_logits parameter in MixtureOfDiagNormals should be K dimensional"
+            ), "The component_logits parameter in MixtureOfDiagNormalsSharedCovariance should be K dimensional"
             assert component_logits.size(0) == locs.size(0)
             batch_shape = ()
         else:
             assert (
                 coord_scale.dim() > 1
-            ), "The coord_scale parameter in MixtureOfDiagNormals should be ... x B x D dimensional"
-            assert (
-                component_logits.dim() > 1
-            ), "The component_logits parameter in MixtureOfDiagNormals should be ... x B x K dimensional"
+            ), "The coord_scale parameter in MixtureOfDiagNormalsSharedCovariance should be ... x B x D dimensional"
+            assert component_logits.dim() > 1, (
+                "The component_logits parameter in MixtureOfDiagNormalsSharedCovariance should be "
+                "... x B x K dimensional"
+            )
             assert component_logits.size(-1) == locs.size(-2)
             batch_shape = tuple(locs.shape[:-2])
         self.locs = locs
@@ -134,7 +137,7 @@ class _MixDiagNormalSharedCovarianceSample(Function):
     @staticmethod
     def forward(ctx, locs, coord_scale, component_logits, pis, which, noise_shape):
         dim = coord_scale.size(-1)
-        white = torch.randn(noise_shape, dtype=locs.dtype, device=locs.device)
+        white = locs.new_empty(noise_shape).normal_()
         n_unsqueezes = locs.dim() - which.dim()
         for _ in range(n_unsqueezes):
             which = which.unsqueeze(-1)


Test patch: