Repository: pyro-ppl/pyro

Issue:
MaskedMixture distribution doesn't work properly when Delta distribution is used as a component
### Issue Description
MaskedMixture fails when Delta distribution is used as a component distribution. In particular, sampling from MaskedMixture outputs wrong values and also alters the Delta distribution used as a component.

### Code Snippet
```py
import sys
import torch
import pyro
import pyro.distributions as dist

print(sys.version)
print(pyro.__version__)
print(torch.__version__)

delta = dist.Delta(torch.tensor([0.]))
gamma = dist.Gamma(torch.ones(2)*100., torch.ones(1))
m = torch.tensor([0, 1]).bool()

print("\nDelta dist before sampling:", delta)

masked_mixture = dist.MaskedMixture(m, delta, gamma)
print("\nSample masked mixture:", masked_mixture.sample())

print("\nDelta dist after sampling:", delta)
```
returns:
```
3.7.6 (default, Jan  8 2020, 19:59:22)
[GCC 7.3.0]
1.2.0
1.4.0

Delta dist before sampling: Delta(v: tensor([0.]), log_density: tensor([0.]))

Sample masked mixture: tensor([103.3137, 103.3137])

Delta dist after sampling: Delta(v: tensor([103.3137]), log_density: tensor([0.]))
```
possible solution (at least it fixes the example above) is to use torch.where in pyro/distributions/mixture.py file like below:

```diff
  def sample(self, sample_shape=torch.Size()):
      mask = self.mask.expand(sample_shape + self.batch_shape) if sample_shape else self.mask
-     result = self.component0.sample(sample_shape)
-     result[mask] = self.component1.sample(sample_shape)[mask]
+     result = torch.where(mask,
+                          self.component1.sample(sample_shape),
+                          self.component0.sample(sample_shape))
      return result
```


Discussion:
@ordabayevy your fix looks great! Do you want to submit that as a PR? You can add a regression test if you want, but I'm ok with just the fix and existing tests.

Resolving patch:
diff --git a/pyro/distributions/mixture.py b/pyro/distributions/mixture.py
--- a/pyro/distributions/mixture.py
+++ b/pyro/distributions/mixture.py
@@ -101,15 +101,19 @@ def expand(self, batch_shape):
             return type(self)(mask, component0, component1)
 
     def sample(self, sample_shape=torch.Size()):
-        mask = self.mask.expand(sample_shape + self.batch_shape) if sample_shape else self.mask
-        result = self.component0.sample(sample_shape)
-        result[mask] = self.component1.sample(sample_shape)[mask]
+        mask = self.mask.reshape(self.mask.shape + (1,) * self.event_dim)
+        mask = mask.expand(sample_shape + self.shape())
+        result = torch.where(mask,
+                             self.component1.sample(sample_shape),
+                             self.component0.sample(sample_shape))
         return result
 
     def rsample(self, sample_shape=torch.Size()):
-        mask = self.mask.expand(sample_shape + self.batch_shape) if sample_shape else self.mask
-        result = self.component0.rsample(sample_shape)
-        result[mask] = self.component1.rsample(sample_shape)[mask]
+        mask = self.mask.reshape(self.mask.shape + (1,) * self.event_dim)
+        mask = mask.expand(sample_shape + self.shape())
+        result = torch.where(mask,
+                             self.component1.rsample(sample_shape),
+                             self.component0.rsample(sample_shape))
         return result
 
     def log_prob(self, value):
@@ -122,8 +126,9 @@ def log_prob(self, value):
         mask = self.mask
         if mask.shape != mask_shape:
             mask = mask.expand(mask_shape)
-        result = self.component0.log_prob(value)
-        result[mask] = self.component1.log_prob(value)[mask]
+        result = torch.where(mask,
+                             self.component1.log_prob(value),
+                             self.component0.log_prob(value))
         return result
 
     @lazy_property


Test patch: