Repository: pyro-ppl/pyro
Pull request: bugfix #2270
Resolves #2270 
Using torch.where seems to resolve the issues related to using the Delta distribution as a component in the MaskedMixture.

Problem statement:
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
```

Hints from issue 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.
@fritzo I submitted a PR with the fix. I am not sure how to do a regression test. This is my first time submitting a bugfix (and contributing) to an open project so please let me know if I am not doing it right!

Commit message: reshape mask (fehiepsi)

Patch to review:
diff --git a/pyro/distributions/mixture.py b/pyro/distributions/mixture.py
index 096c8137c2..36c46bc0be 100644
--- a/pyro/distributions/mixture.py
+++ b/pyro/distributions/mixture.py
@@ -102,14 +102,17 @@ def expand(self, batch_shape):
 
     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 = mask.reshape(mask.shape + (1,) * len(self.event_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]
+        result = torch.where(mask,
+                             self.component1.rsample(sample_shape),
+                             self.component0.rsample(sample_shape))
         return result
 
     def log_prob(self, value):
@@ -122,8 +125,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


Code review comments:
diff_hunk: @@ -102,14 +102,17 @@ def expand(self, batch_shape):
 
     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 = mask.reshape(mask.shape + (1,) * len(self.event_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]
path: pyro/distributions/mixture.py
text: @user1:
@author Could you also reshape the mask here?

@author:
@user1 ok, thanks. Reshaping in the log_prob seems to be unnecessary.

@user1:
I think it is required. There, we also want the `mask` shape to include `event_shape`.

@author:
including it there fails the test, presumably because the output of the log_prob has the shape of `sample_shape + batch_shape` without the `event_shape`

@user1:
Opps, you are right!!!