Repository: pytorch/pytorch
Pull request: Implement OneHotCategorical distribution
Fixes probtorch#5 (OneHotCategorical)
Fixes probtorch#8 (.shape properties)
Fixes #4353 (fixing random seeds in tests)

This implements a `OneHotCategorial` distribution as a thin wrapper around `Categorical`. This is the most well-tested multivariate distribution so far (`Dirichlet` is the only other multivariate distribution and that is mostly tested by the `Beta` special case). I've fixed a number of issues with `test_distributions.py` to allow easier testing of multivariate distributions.

This also adds `.batch_shape` and `.event_shape` properties, which are now used by `OneHotCategorical` and in tests.

Finally this sets random number seeds in `common.TestCase.setUp()` method as suggested by @ezyang. I've added that in this PR because a nondeterministic test failure came up while testing this locally. This PR also adds maintainer instructions for dealing with randomized statistical tests, as suggested by @ezyang.

Math was reviewed by @neerajprad 

Problem statement:
Implement OneHotCategorical


Add .shape properties to all distributions


test_distributions.py should seed on setUp, to make it harder to accidentally commit nondeterministic tests
Right now it seeds only if you call `_set_rng_seed` which means it's easy to forget to call it, leading to a nondeterministic test.

Instead, we should modify `setUp` in `common.py` and have it seed Numpy as well as Torch:

```
class TestCase(unittest.TestCase):
    precision = 1e-5
    maxDiff = None

    def setUp(self):
        torch.manual_seed(SEED)
```

CC @fritzo

Hints from issue discussion:
TensorFlow defines [tf.contrib.distributions.OneHotCategorical](https://www.tensorflow.org/api_docs/python/tf/contrib/distributions/OneHotCategorical) for this purpose, which is a distinct class from [tf.contrib.distributions.Categorical](https://www.tensorflow.org/api_docs/python/tf/contrib/distributions/Categorical). I think this is probably a good idea, since samples from the two distributions have different types (one-hot vector, vs integer).
Should this be called `shape` or `size`? On the one hand, `size` is more similar to what PyTorch already provides for Tensors. On the other hand, `size()` is a method, not a property. 
`shape` seems preferable since (1) PyTorch is moving towards the numpy api (already has partial support for shape), and (2) edward uses shape so code will be more interchangeable.
We should also add some maintainer-instructions for dealing with changing randomness sources in test_distributions.py. In particular there are two types of randomized tests:
1. The easier type of randomized test are tests that should always pass but are initialized with random data. If these fail something is wrong, but it's fine to use a fixed seed.
2. The trickier tests are statistical tests, and these are explicitly marked "this is a randomized test". These have a known positive failure rate (we set `failure_rate=1e-3` by default). We need to balance strength of these tests with annoyance of false alarms. One way that works is to specifically set seeds in each of the randomized tests. When a random generator occasionally changes (as in #4312 ), some of these statistical tests may (rarely) fail. If one fails in this case, it's fine to increment the seed of the failing test (but you shouldn't need to increment it more than once; otherwise something is probably actually wrong).

@ezyang If you add a generic `torch.manual_seed(0)` in a `setUp()`, please do not modify or remove `_set_rng_seed()` calls in any tests marked "this is a randomized test", since these statistical tests need individual seeds.
I'll try to add a `setUp` method as part of the `OneHotCategorical` PR (the nondeterminism just bit me 😄 )

Commit message: Flake8

Patch to review:
diff --git a/docs/source/distributions.rst b/docs/source/distributions.rst
index 5843cf988a59..d0c7fd7243e2 100644
--- a/docs/source/distributions.rst
+++ b/docs/source/distributions.rst
@@ -48,3 +48,15 @@ Probability distributions - torch.distributions
 
 .. autoclass:: Normal
     :members:
+
+:hidden:`OneHotCategorical`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: OneHotCategorical
+    :members:
+
+:hidden:`Uniform`
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: Uniform
+    :members:
diff --git a/test/test_distributions.py b/test/test_distributions.py
index 776cc60a274e..dd1048d1830a 100644
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -7,8 +7,8 @@
 from common import TestCase, run_tests
 from torch.autograd import Variable, gradcheck
 from torch.distributions import (Bernoulli, Beta, Categorical, Dirichlet,
-                                 Exponential, Gamma, Laplace, Normal)
-from torch.distributions.uniform import Uniform
+                                 Exponential, Gamma, Laplace, Normal,
+                                 OneHotCategorical, Uniform)
 
 TEST_NUMPY = True
 try:
@@ -19,6 +19,12 @@
     TEST_NUMPY = False
 
 
+def set_rng_seed(seed=0):
+    torch.manual_seed(seed)
+    if TEST_NUMPY:
+        np.random.seed(seed)
+
+
 # Register all distributions for generic tests.
 Example = namedtuple('Example', ['Dist', 'params'])
 EXAMPLES = [
@@ -41,6 +47,10 @@
         {'probs': Variable(torch.Tensor([[0.1, 0.2, 0.3], [0.5, 0.3, 0.2]]), requires_grad=True)},
         {'probs': Variable(torch.Tensor([[1.0, 0.0], [0.0, 1.0]]), requires_grad=True)},
     ]),
+    Example(OneHotCategorical, [
+        {'probs': Variable(torch.Tensor([[0.1, 0.2, 0.3], [0.5, 0.3, 0.2]]), requires_grad=True)},
+        {'probs': Variable(torch.Tensor([[1.0, 0.0], [0.0, 1.0]]), requires_grad=True)},
+    ]),
     Example(Gamma, [
         {
             'alpha': Variable(torch.exp(torch.randn(2, 3)), requires_grad=True),
@@ -105,17 +115,18 @@
 
 
 class TestDistributions(TestCase):
-    def _set_rng_seed(self, seed=0):
-        torch.manual_seed(seed)
-        if TEST_NUMPY:
-            np.random.seed(seed)
+    def setUp(self):
+        set_rng_seed(0)
 
     def _gradcheck_log_prob(self, dist_ctor, ctor_params):
         # performs gradient checks on log_prob
         distribution = dist_ctor(*ctor_params)
         s = distribution.sample()
 
-        self.assertEqual(s.size(), distribution.log_prob(s).size())
+        expected_shape = distribution.batch_shape + distribution.event_shape
+        if not expected_shape:
+            expected_shape = torch.Size((1,))  # Work around lack of scalars.
+        self.assertEqual(s.size(), expected_shape)
 
         def apply_fn(*params):
             return dist_ctor(*params).log_prob(s)
@@ -185,10 +196,7 @@ def ref_log_prob(idx, val, log_prob):
             self.assertEqual(log_prob, math.log(prob if val else 1 - prob))
 
         self._check_log_prob(Bernoulli(p), ref_log_prob)
-
-        def call_rsample():
-            return Bernoulli(r).rsample()
-        self.assertRaises(NotImplementedError, call_rsample)
+        self.assertRaises(NotImplementedError, Bernoulli(r).rsample)
 
     def test_bernoulli_enumerate_support(self):
         examples = [
@@ -212,10 +220,7 @@ def test_categorical_1d(self):
         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,))
-
-        def call_rsample():
-            return Categorical(p).rsample()
-        self.assertRaises(NotImplementedError, call_rsample)
+        self.assertRaises(NotImplementedError, Categorical(p).rsample)
 
     def test_categorical_2d(self):
         probabilities = [[0.1, 0.2, 0.3], [0.5, 0.3, 0.2]]
@@ -228,7 +233,7 @@ def test_categorical_2d(self):
         self._gradcheck_log_prob(Categorical, (p,))
 
         # sample check for extreme value of probs
-        self._set_rng_seed(0)
+        set_rng_seed(0)
         self.assertEqual(Categorical(s).sample(sample_shape=(2,)).data,
                          torch.Tensor([[0, 1], [0, 1]]))
 
@@ -245,6 +250,35 @@ def test_categorical_enumerate_support(self):
         ]
         self._check_enumerate_support(Categorical, examples)
 
+    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.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,))
+        self.assertRaises(NotImplementedError, OneHotCategorical(p).rsample)
+
+    def test_one_hot_categorical_2d(self):
+        probabilities = [[0.1, 0.2, 0.3], [0.5, 0.3, 0.2]]
+        probabilities_1 = [[1.0, 0.0], [0.0, 1.0]]
+        p = Variable(torch.Tensor(probabilities), requires_grad=True)
+        s = Variable(torch.Tensor(probabilities_1), requires_grad=True)
+        self.assertEqual(OneHotCategorical(p).sample().size(), (2, 3))
+        self.assertEqual(OneHotCategorical(p).sample(sample_shape=(3, 4)).size(), (3, 4, 2, 3))
+        self.assertEqual(OneHotCategorical(p).sample_n(6).size(), (6, 2, 3))
+        self._gradcheck_log_prob(OneHotCategorical, (p,))
+
+        dist = OneHotCategorical(p)
+        x = dist.sample()
+        self.assertEqual(dist.log_prob(x), Categorical(p).log_prob(x.max(-1)[1]))
+
+    def test_one_hot_categorical_enumerate_support(self):
+        examples = [
+            ([0.1, 0.2, 0.7], [[1, 0, 0], [0, 1, 0], [0, 0, 1]]),
+            ([[0.1, 0.9], [0.3, 0.7]], [[[1, 0], [1, 0]], [[0, 1], [0, 1]]]),
+        ]
+        self._check_enumerate_support(OneHotCategorical, examples)
+
     def test_uniform(self):
         low = Variable(torch.zeros(5, 5), requires_grad=True)
         high = Variable(torch.ones(5, 5) * 3, requires_grad=True)
@@ -263,7 +297,7 @@ def test_uniform(self):
         self.assertEqual(uniform.log_prob(above_high).data[0], -float('inf'), allow_inf=True)
         self.assertEqual(uniform.log_prob(below_low).data[0], -float('inf'), allow_inf=True)
 
-        self._set_rng_seed(1)
+        set_rng_seed(1)
         self._gradcheck_log_prob(Uniform, (low, high))
         self._gradcheck_log_prob(Uniform, (low, 1.0))
         self._gradcheck_log_prob(Uniform, (0.0, high))
@@ -293,7 +327,7 @@ def test_normal(self):
         self.assertEqual(Normal(-0.7, 50.0).sample_n(1).size(), (1,))
 
         # sample check for extreme value of mean, std
-        self._set_rng_seed(1)
+        set_rng_seed(1)
         self.assertEqual(Normal(mean_delta, std_delta).sample(sample_shape=(1, 2)),
                          torch.Tensor([[[1.0, 0.0], [1.0, 0.0]]]),
                          prec=1e-4)
@@ -325,7 +359,7 @@ def ref_log_prob(idx, x, log_prob):
     # This is a randomized test.
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_normal_sample(self):
-        self._set_rng_seed()
+        set_rng_seed()
         for mean, std in product([-1.0, 0.0, 1.0], [0.1, 1.0, 10.0]):
             self._check_sampler_sampler(Normal(mean, std),
                                         scipy.stats.norm(loc=mean, scale=std),
@@ -361,7 +395,7 @@ def ref_log_prob(idx, x, log_prob):
     # This is a randomized test.
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_exponential_sample(self):
-        self._set_rng_seed(1)
+        set_rng_seed(1)
         for rate in [1e-5, 1.0, 10.]:
             self._check_sampler_sampler(Exponential(rate),
                                         scipy.stats.expon(scale=1. / rate),
@@ -382,7 +416,7 @@ def test_laplace(self):
         self.assertEqual(Laplace(-0.7, 50.0).sample_n(1).size(), (1,))
 
         # sample check for extreme value of mean, std
-        self._set_rng_seed()
+        set_rng_seed()
         self.assertEqual(Laplace(loc_delta, scale_delta).sample(sample_shape=(1, 2)),
                          torch.Tensor([[[1.0, 0.0], [1.0, 0.0]]]),
                          prec=1e-4)
@@ -413,13 +447,12 @@ def ref_log_prob(idx, x, log_prob):
     # This is a randomized test.
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_laplace_sample(self):
-        self._set_rng_seed(1)
+        set_rng_seed(1)
         for loc, scale in product([-1.0, 0.0, 1.0], [0.1, 1.0, 10.0]):
             self._check_sampler_sampler(Laplace(loc, scale),
                                         scipy.stats.laplace(loc=loc, scale=scale),
                                         'Laplace(loc={}, scale={})'.format(loc, scale))
 
-    # This is a randomized test.
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_gamma_shape(self):
         alpha = Variable(torch.exp(torch.randn(2, 3)), requires_grad=True)
@@ -444,7 +477,7 @@ def ref_log_prob(idx, x, log_prob):
     # This is a randomized test.
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_gamma_sample(self):
-        self._set_rng_seed()
+        set_rng_seed()
         for alpha, beta in product([0.1, 1.0, 5.0], [0.1, 1.0, 10.0]):
             self._check_sampler_sampler(Gamma(alpha, beta),
                                         scipy.stats.gamma(alpha, scale=1.0 / beta),
@@ -453,7 +486,7 @@ def test_gamma_sample(self):
     # This is a randomized test.
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_gamma_sample_grad(self):
-        self._set_rng_seed(1)
+        set_rng_seed(1)
         num_samples = 100
         for alpha in [1e-3, 1e-2, 1e-1, 1e0, 1e1, 1e2, 1e3, 1e4]:
             alphas = Variable(torch.Tensor([alpha] * num_samples), requires_grad=True)
@@ -501,7 +534,7 @@ def test_dirichlet_log_prob(self):
     # This is a randomized test.
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_dirichlet_sample(self):
-        self._set_rng_seed()
+        set_rng_seed()
         alpha = torch.exp(torch.randn(3))
         self._check_sampler_sampler(Dirichlet(alpha),
                                     scipy.stats.dirichlet(alpha.numpy()),
@@ -534,7 +567,7 @@ def test_beta_log_prob(self):
     # This is a randomized test.
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_beta_sample(self):
-        self._set_rng_seed(1)
+        set_rng_seed(1)
         for alpha, beta in product([0.1, 1.0, 10.0], [0.1, 1.0, 10.0]):
             self._check_sampler_sampler(Beta(alpha, beta),
                                         scipy.stats.beta(alpha, beta),
@@ -547,7 +580,7 @@ def test_beta_sample(self):
     # This is a randomized test.
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_beta_sample_grad(self):
-        self._set_rng_seed()
+        set_rng_seed()
         num_samples = 20
         for alpha, beta in product([1e-2, 1e0, 1e2], [1e-2, 1e0, 1e2]):
             alphas = Variable(torch.Tensor([alpha] * num_samples), requires_grad=True)
@@ -649,6 +682,7 @@ def test_invalid_parameter_broadcasting(self):
 
 class TestDistributionShapes(TestCase):
     def setUp(self):
+        set_rng_seed(0)
         self.scalar_sample = 1
         self.tensor_sample_1 = torch.ones(3, 2)
         self.tensor_sample_2 = torch.ones(3, 2, 3)
@@ -705,13 +739,23 @@ def test_beta_shape_tensor_params(self):
         self.assertRaises(ValueError, dist.log_prob, self.tensor_sample_2)
 
     def test_categorical_shape(self):
-        categorical = Categorical(torch.Tensor([[0.6, 0.3], [0.6, 0.3], [0.6, 0.3]]))
-        self.assertEqual(categorical._batch_shape, torch.Size((3,)))
-        self.assertEqual(categorical._event_shape, torch.Size(()))
-        self.assertEqual(categorical.sample().size(), torch.Size((3,)))
-        self.assertEqual(categorical.sample((3, 2)).size(), torch.Size((3, 2, 3,)))
-        self.assertRaises(ValueError, categorical.log_prob, self.tensor_sample_1)
-        self.assertEqual(categorical.log_prob(self.tensor_sample_2).size(), torch.Size((3, 2, 3)))
+        dist = Categorical(torch.Tensor([[0.6, 0.3], [0.6, 0.3], [0.6, 0.3]]))
+        self.assertEqual(dist._batch_shape, torch.Size((3,)))
+        self.assertEqual(dist._event_shape, torch.Size(()))
+        self.assertEqual(dist.sample().size(), torch.Size((3,)))
+        self.assertEqual(dist.sample((3, 2)).size(), torch.Size((3, 2, 3,)))
+        self.assertRaises(ValueError, dist.log_prob, self.tensor_sample_1)
+        self.assertEqual(dist.log_prob(self.tensor_sample_2).size(), torch.Size((3, 2, 3)))
+
+    def test_one_hot_categorical_shape(self):
+        dist = OneHotCategorical(torch.Tensor([[0.6, 0.3], [0.6, 0.3], [0.6, 0.3]]))
+        self.assertEqual(dist._batch_shape, torch.Size((3,)))
+        self.assertEqual(dist._event_shape, torch.Size((2,)))
+        self.assertEqual(dist.sample().size(), torch.Size((3, 2)))
+        self.assertEqual(dist.sample((3, 2)).size(), torch.Size((3, 2, 3, 2)))
+        self.assertEqual(dist.log_prob(self.tensor_sample_1).size(), torch.Size((3,)))
+        self.assertRaises(ValueError, dist.log_prob, self.tensor_sample_2)
+        self.assertEqual(dist.log_prob(dist.enumerate_support()).size(), torch.Size((2, 3)))
 
     def test_dirichlet_shape(self):
         dist = Dirichlet(torch.Tensor([[0.6, 0.3], [1.6, 1.3], [2.6, 2.3]]))
diff --git a/torch/distributions/__init__.py b/torch/distributions/__init__.py
index 3196e37c32fc..0035cdfa15f5 100644
--- a/torch/distributions/__init__.py
+++ b/torch/distributions/__init__.py
@@ -35,10 +35,23 @@
 from .categorical import Categorical
 from .dirichlet import Dirichlet
 from .distribution import Distribution
-from .gamma import Gamma
-from .normal import Normal
 from .exponential import Exponential
+from .gamma import Gamma
 from .laplace import Laplace
-
-
-__all__ = ['Distribution', 'Bernoulli', 'Beta', 'Categorical', 'Dirichlet', 'Gamma', 'Normal', 'Exponential', 'Laplace']
+from .normal import Normal
+from .one_hot_categorical import OneHotCategorical
+from .uniform import Uniform
+
+__all__ = [
+    'Bernoulli',
+    'Beta',
+    'Categorical',
+    'Dirichlet',
+    'Distribution',
+    'Exponential',
+    'Gamma',
+    'Laplace',
+    'Normal',
+    'OneHotCategorical',
+    'Uniform',
+]
diff --git a/torch/distributions/distribution.py b/torch/distributions/distribution.py
index c0af79873550..58ed0054a918 100644
--- a/torch/distributions/distribution.py
+++ b/torch/distributions/distribution.py
@@ -15,6 +15,20 @@ def __init__(self, batch_shape=torch.Size(), event_shape=torch.Size()):
         self._batch_shape = batch_shape
         self._event_shape = event_shape
 
+    @property
+    def batch_shape(self):
+        """
+        Returns the shape over which parameters are batched.
+        """
+        return self._batch_shape
+
+    @property
+    def event_shape(self):
+        """
+        Returns the shape of a single sample (without batching).
+        """
+        return self._event_shape
+
     def sample(self, sample_shape=torch.Size()):
         """
         Generates a sample_shape shaped sample or sample_shape shaped batch of
@@ -74,7 +88,7 @@ def entropy(self):
         """
         raise NotImplementedError
 
-    def _extended_shape(self, sample_shape=()):
+    def _extended_shape(self, sample_shape=torch.Size()):
         """
         Returns the size of the sample returned by the distribution, given
         a `sample_shape`. Note, that the batch and event shapes of a distribution
@@ -84,7 +98,7 @@ def _extended_shape(self, sample_shape=()):
         Args:
             sample_shape (torch.Size): the size of the sample to be drawn.
         """
-        shape = sample_shape + self._batch_shape + self._event_shape
+        shape = torch.Size(sample_shape + self._batch_shape + self._event_shape)
         if not shape:
             shape = torch.Size((1,))
         return shape
diff --git a/torch/distributions/one_hot_categorical.py b/torch/distributions/one_hot_categorical.py
new file mode 100644
index 000000000000..514b157bbb3b
--- /dev/null
+++ b/torch/distributions/one_hot_categorical.py
@@ -0,0 +1,60 @@
+import torch
+from torch.autograd import Variable
+from torch.distributions.distribution import Distribution
+from torch.distributions.categorical import Categorical
+
+
+class OneHotCategorical(Distribution):
+    r"""
+    Creates a one-hot categorical distribution parameterized by `probs`.
+
+    Samples are one-hot coded vectors of size probs.size(-1).
+
+    See also: :func:`torch.distributions.Categorical`
+
+    Example::
+
+        >>> m = OneHotCategorical(torch.Tensor([ 0.25, 0.25, 0.25, 0.25 ]))
+        >>> m.sample()  # equal probability of 0, 1, 2, 3
+         0
+         0
+         1
+         0
+        [torch.FloatTensor of size 4]
+
+    Args:
+        probs (Tensor or Variable): event probabilities
+    """
+    has_enumerate_support = True
+
+    def __init__(self, probs):
+        self._categorical = Categorical(probs)
+        batch_shape = probs.size()[:-1]
+        event_shape = probs.size()[-1:]
+        super(OneHotCategorical, self).__init__(batch_shape, event_shape)
+
+    def sample(self, sample_shape=torch.Size()):
+        sample_shape = torch.Size(sample_shape)
+        probs = self._categorical.probs
+        one_hot = probs.new(self._extended_shape(sample_shape)).zero_()
+        indices = self._categorical.sample(sample_shape)
+        if indices.dim() < one_hot.dim():
+            indices = indices.unsqueeze(-1)
+        return one_hot.scatter_(-1, indices, 1)
+
+    def log_prob(self, value):
+        indices = value.max(-1)[1]
+        return self._categorical.log_prob(indices)
+
+    def entropy(self):
+        return self._categorical.entropy()
+
+    def enumerate_support(self):
+        probs = self._categorical.probs
+        n = self.event_shape[0]
+        if isinstance(probs, Variable):
+            values = Variable(torch.eye(n, out=probs.data.new(n, n)))
+        else:
+            values = torch.eye(n, out=probs.new(n, n))
+        values = values.view((n,) + (1,) * len(self.batch_shape) + (n,))
+        return values.expand((n,) + self.batch_shape + (n,))


Code review comments:
diff_hunk: @@ -105,17 +115,18 @@
 
 
 class TestDistributions(TestCase):
-    def _set_rng_seed(self, seed=0):
-        torch.manual_seed(seed)
-        if TEST_NUMPY:
-            np.random.seed(seed)
+    def setUp(self):
+        set_rng_seed(0)
path: test/test_distributions.py
text: @user1:
I think I'd prefer it if the numpy seeding was moved to `test/common.py` (with appropriate conditionals for whether or not numpy is available); we have numpy tests that are not in `test_distributions.py`, but if the random seeding is done here then no other files can benefit. There is an added hazard here which is that if `setUp` in `TestCase` ever changes, the fact that we are not calling the superclass method means that `TestDistributions` will silently fail to get this behavior.

@author:
Good points! I've moved `.setUp()` to the base class and moved `set_rng_seed()` to tests/common.py.