Repository: pytorch/pytorch

Issue:
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


Discussion:
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 😄 )

Resolving patch:
diff --git a/torch/distributions/__init__.py b/torch/distributions/__init__.py
--- 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
--- 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
--- /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,))


Test patch:
diff --git a/test/common.py b/test/common.py
--- a/test/common.py
+++ b/test/common.py
@@ -97,6 +97,12 @@ def to_gpu(obj, type_map={}):
         return deepcopy(obj)
 
 
+def set_rng_seed(seed):
+    torch.manual_seed(seed)
+    if TEST_NUMPY:
+        numpy.random.seed(seed)
+
+
 @contextlib.contextmanager
 def freeze_rng_state():
     rng_state = torch.get_rng_state()
@@ -129,7 +135,7 @@ class TestCase(unittest.TestCase):
     maxDiff = None
 
     def setUp(self):
-        torch.manual_seed(SEED)
+        set_rng_seed(SEED)
 
     def assertTensorsSlowEqual(self, x, y, prec=None, message=''):
         max_err = 0
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -1,14 +1,38 @@
+"""
+Note [Randomized statistical tests]
+-----------------------------------
+
+This note describes how to maintain tests in this file as random sources
+change. This file contains 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 by inheriting from common.TestCase.
+
+2. The trickier tests are statistical tests. These tests explicitly call
+   set_rng_seed(n) and are marked "see Note [Randomized statistical tests]".
+   These statistical tests 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 vectorizing the Box-Muller sampler), 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).
+"""
+
 import math
 import unittest
 from collections import namedtuple
 from itertools import product
 
 import torch
-from common import TestCase, run_tests
+from common import TestCase, run_tests, set_rng_seed
 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:
@@ -41,6 +65,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 +133,15 @@
 
 
 class TestDistributions(TestCase):
-    def _set_rng_seed(self, seed=0):
-        torch.manual_seed(seed)
-        if TEST_NUMPY:
-            np.random.seed(seed)
-
     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 +211,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 +235,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 +248,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 +265,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 +312,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 +342,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)
@@ -322,10 +371,9 @@ def ref_log_prob(idx, x, log_prob):
 
         self._check_log_prob(Normal(mean, std), ref_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(0)  # see Note [Randomized statistical tests]
         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),
@@ -358,10 +406,9 @@ def ref_log_prob(idx, x, log_prob):
 
         self._check_log_prob(Exponential(rate), ref_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)  # see Note [Randomized statistical tests]
         for rate in [1e-5, 1.0, 10.]:
             self._check_sampler_sampler(Exponential(rate),
                                         scipy.stats.expon(scale=1. / rate),
@@ -382,7 +429,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(0)
         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)
@@ -410,16 +457,14 @@ def ref_log_prob(idx, x, log_prob):
 
         self._check_log_prob(Laplace(loc, scale), ref_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)  # see Note [Randomized statistical tests]
         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)
@@ -441,19 +486,17 @@ def ref_log_prob(idx, x, log_prob):
 
         self._check_log_prob(Gamma(alpha, beta), ref_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(0)  # see Note [Randomized statistical tests]
         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),
                                         'Gamma(alpha={}, beta={})'.format(alpha, beta))
 
-    # 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)  # see Note [Randomized statistical tests]
         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)
@@ -498,10 +541,9 @@ def test_dirichlet_log_prob(self):
             expected_log_prob = scipy.stats.dirichlet.logpdf(x[i].numpy(), alpha.numpy())
             self.assertAlmostEqual(actual_log_prob[i], expected_log_prob, places=3)
 
-    # 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(0)  # see Note [Randomized statistical tests]
         alpha = torch.exp(torch.randn(3))
         self._check_sampler_sampler(Dirichlet(alpha),
                                     scipy.stats.dirichlet(alpha.numpy()),
@@ -531,10 +573,9 @@ def test_beta_log_prob(self):
             expected_log_prob = scipy.stats.beta.logpdf(x, alpha, beta)[0]
             self.assertAlmostEqual(actual_log_prob, expected_log_prob, places=3)
 
-    # 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)  # see Note [Randomized statistical tests]
         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),
@@ -544,10 +585,9 @@ def test_beta_sample(self):
             x = Beta(Tensor([1e-6]), Tensor([1e-6])).sample()[0]
             self.assertTrue(np.isfinite(x) and x > 0, 'Invalid Beta.sample(): {}'.format(x))
 
-    # 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(0)  # see Note [Randomized statistical tests]
         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 +689,7 @@ def test_invalid_parameter_broadcasting(self):
 
 class TestDistributionShapes(TestCase):
     def setUp(self):
+        super(TestCase, self).setUp()
         self.scalar_sample = 1
         self.tensor_sample_1 = torch.ones(3, 2)
         self.tensor_sample_2 = torch.ones(3, 2, 3)
@@ -705,13 +746,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]]))