Repository: pytorch/pytorch
Pull request: Fix return type for Bernoulli enumerate_support
Fixes [probtorch/pytorch#68](https://github.com/probtorch/pytorch/issues/68). Only the return type for Bernoulli's `enumerate_support` needed change. 
 - We use `FloatTensor` as the default return type of distribution's samples, , as most tensor operations support `FloatTensor` without needing any casting. 
 - If the underlying torch sampler has a different tensor return type, that may be used instead. e.g. in the case of Categorical, we use `torch.multinomial` that returns a `LongTensor`. The return type of `Categorical` is therefore a `LongTensor`. This also helps in cases where the generated sample is used to index into tensors or numpy arrays. However, for the `OneHotCategorical`, we default to `FloatTensor` since the samples are more generally used as masks.

Problem statement:
Ensure types are consistent and check types in tests
Currently the tensor type situation is a mess.
- `Bernoulli.sample()` returns a `FloatTensor`
- `Bernoulli.enumerate_support()` returns a `LongTensor`
- `Categorical.sample()` returns a `LongTensor`
- `Categorical.enumerate_support()` returns a `LongTensor`
- Everything else returns a `FloatTensor`

I believe the right thing to do is:
1. ensure `.enumerate_support()` type matches `.sample()` type.
2. try to match `torch.*` methods like `torch.bernoulli` and `torch.multinomial`.

We should also check that types agree in tests. This can be a simple parameterized test like
```py
for Dist, params in EXAMPLES:
    for i, param in numerate(params):
        dist = Dist(**params)
        self.assertTrue(type(dist.sample()) is type(dist.enumerate_support()))
```

Hints from issue discussion:

Commit message: Fix return type for Bernoulli support

Patch to review:
diff --git a/test/test_distributions.py b/test/test_distributions.py
index 75f45d63068b..30ef6c6e02fa 100644
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -163,6 +163,12 @@
 ]
 
 
+def unwrap(value):
+    if isinstance(value, Variable):
+        return value.data
+    return value
+
+
 class TestDistributions(TestCase):
     def _gradcheck_log_prob(self, dist_ctor, ctor_params):
         # performs gradient checks on log_prob
@@ -226,11 +232,23 @@ def _check_enumerate_support(self, dist, examples):
             actual = dist(param).enumerate_support()
             self.assertEqual(actual, expected)
 
+    def test_support_return_type(self):
+        for Dist, params in EXAMPLES:
+            for i, param in enumerate(params):
+                dist = Dist(**param)
+                try:
+                    self.assertTrue(type(unwrap(dist.sample())) is type(unwrap(dist.enumerate_support())),
+                                    msg=('{} example {}/{}, return type mismatch between ' +
+                                         'sample and enumerate_support.').format(Dist.__name__, i, len(params)))
+                except NotImplementedError:
+                    pass
+
     def test_bernoulli(self):
         p = Variable(torch.Tensor([0.7, 0.2, 0.4]), requires_grad=True)
         r = Variable(torch.Tensor([0.3]), requires_grad=True)
         s = 0.3
         self.assertEqual(Bernoulli(p).sample_n(8).size(), (8, 3))
+        self.assertTrue(isinstance(Bernoulli(p).sample().data, torch.Tensor))
         self.assertEqual(Bernoulli(r).sample_n(8).size(), (8, 1))
         self.assertEqual(Bernoulli(r).sample().size(), (1,))
         self.assertEqual(Bernoulli(r).sample((3, 2)).size(), (3, 2, 1))
@@ -269,6 +287,7 @@ def test_categorical_1d(self):
         p = Variable(torch.Tensor([0.1, 0.2, 0.3]), requires_grad=True)
         # TODO: this should return a 0-dim tensor once we have Scalar support
         self.assertEqual(Categorical(p).sample().size(), (1,))
+        self.assertTrue(isinstance(Categorical(p).sample().data, torch.LongTensor))
         self.assertEqual(Categorical(p).sample((2, 2)).size(), (2, 2))
         self.assertEqual(Categorical(p).sample_n(1).size(), (1,))
         self._gradcheck_log_prob(Categorical, (p,))
@@ -310,6 +329,7 @@ def test_categorical_enumerate_support(self):
     def test_one_hot_categorical_1d(self):
         p = Variable(torch.Tensor([0.1, 0.2, 0.3]), requires_grad=True)
         self.assertEqual(OneHotCategorical(p).sample().size(), (3,))
+        self.assertTrue(isinstance(OneHotCategorical(p).sample().data, torch.Tensor))
         self.assertEqual(OneHotCategorical(p).sample((2, 2)).size(), (2, 2, 3))
         self.assertEqual(OneHotCategorical(p).sample_n(1).size(), (1, 3))
         self._gradcheck_log_prob(OneHotCategorical, (p,))
diff --git a/torch/distributions/bernoulli.py b/torch/distributions/bernoulli.py
index d3905001c971..a5607245582e 100644
--- a/torch/distributions/bernoulli.py
+++ b/torch/distributions/bernoulli.py
@@ -64,7 +64,7 @@ def entropy(self):
         return binary_cross_entropy_with_logits(self.logits, self.probs, reduce=False)
 
     def enumerate_support(self):
-        values = torch.arange(2).long()
+        values = torch.arange(2)
         values = values.view((-1,) + (1,) * len(self._batch_shape))
         values = values.expand((-1,) + self._batch_shape)
         if self.probs.is_cuda:


Code review comments:
diff_hunk: @@ -226,11 +232,23 @@ def _check_enumerate_support(self, dist, examples):
             actual = dist(param).enumerate_support()
             self.assertEqual(actual, expected)
 
+    def test_support_return_type(self):
path: test/test_distributions.py
text: @user1:
nit: Maybe call this `test_enumerate_support_type` since `.support` is a property with different meaning?

@author:
Done.