Fix distributions.Categorical.sample bug from .view() (#23328)

fritzo · facebook-github-bot · commit 06762b47212d · 2019-07-29T12:09:50.000-07:00
Summary: This modernizes distributions code by replacing a few uses of `.contiguous().view()` with `.reshape()`, fixing a sample bug in the `Categorical` distribution. The bug is exercised by the following test: ```py batch_shape = (1, 2, 1, 3, 1) sample_shape = (4,) cardinality = 2 logits = torch.randn(batch_shape + (cardinality,)) dist.Categorical(logits=logits).sample(sample_shape) # RuntimeError: invalid argument 2: view size is not compatible with # input tensor's size and stride (at least one dimension spans across # two contiguous subspaces). Call .contiguous() before .view(). # at ../aten/src/TH/generic/THTensor.cpp:203 ``` I have verified this works locally, but I have not added this as a regression test because it is unlikely to regress (the code is now simpler). Pull Request resolved: #23328 Differential Revision: D16510678 Pulled By: colesbury fbshipit-source-id: c125c1a37d21d185132e8e8b65241c86ad8ad04b
diff --git a/torch/distributions/categorical.py b/torch/distributions/categorical.py
@@ -103,12 +103,9 @@ def sample(self, sample_shape=torch.Size()):
         sample_shape = self._extended_shape(sample_shape)
         param_shape = sample_shape + torch.Size((self._num_events,))
         probs = self.probs.expand(param_shape)
-        if self.probs.dim() == 1 or self.probs.size(0) == 1:
-            probs_2d = probs.view(-1, self._num_events)
-        else:
-            probs_2d = probs.contiguous().view(-1, self._num_events)
+        probs_2d = probs.reshape(-1, self._num_events)
         sample_2d = torch.multinomial(probs_2d, 1, True)
-        return sample_2d.contiguous().view(sample_shape)
+        return sample_2d.reshape(sample_shape)
 
     def log_prob(self, value):
         if self._validate_args:
diff --git a/torch/distributions/transforms.py b/torch/distributions/transforms.py
@@ -555,7 +555,7 @@ def _call(self, x):
         return torch.stack([self._call_on_event(flat_x[i]) for i in range(flat_x.size(0))]).view(x.shape)
 
     def _inverse(self, y):
-        flat_y = y.contiguous().view((-1,) + y.shape[-2:])
+        flat_y = y.reshape((-1,) + y.shape[-2:])
         return torch.stack([self._inverse_on_event(flat_y[i]) for i in range(flat_y.size(0))]).view(y.shape)
 
 
