Merge pull request #134 from iotamudelta/distri_tests

Enable distribution tests

Author: iotamudelta

test/run_test.py

@@ -46,7 +46,6 @@
     'c10d',
     'cpp_extensions',
     'distributed',
-    'distributions',
     'multiprocessing',
     'nccl',
     'thd_distributed',

test/test_distributions.py

@@ -31,7 +31,7 @@
 
 import torch
 from torch._six import inf
-from common import TestCase, run_tests, set_rng_seed, TEST_WITH_UBSAN
+from common import TestCase, run_tests, set_rng_seed, TEST_WITH_UBSAN, skipIfRocm
 from common_cuda import TEST_CUDA
 from torch.autograd import grad, gradcheck
 from torch.distributions import (Bernoulli, Beta, Binomial, Categorical,
@@ -724,6 +724,7 @@ def test_repr(self):
                 dist = Dist(**param)
                 self.assertTrue(repr(dist).startswith(dist.__class__.__name__))
 
+    @skipIfRocm
     def test_sample_detached(self):
         for Dist, params in EXAMPLES:
             for i, param in enumerate(params):
@@ -736,6 +737,7 @@ def test_sample_detached(self):
                                  msg='{} example {}/{}, .sample() is not detached'.format(
                                      Dist.__name__, i + 1, len(params)))
 
+    @skipIfRocm
     def test_rsample_requires_grad(self):
         for Dist, params in EXAMPLES:
             for i, param in enumerate(params):
@@ -749,6 +751,7 @@ def test_rsample_requires_grad(self):
                                 msg='{} example {}/{}, .rsample() does not require grad'.format(
                                     Dist.__name__, i + 1, len(params)))
 
+    @skipIfRocm
     def test_enumerate_support_type(self):
         for Dist, params in EXAMPLES:
             for i, param in enumerate(params):
@@ -760,6 +763,7 @@ def test_enumerate_support_type(self):
                 except NotImplementedError:
                     pass
 
+    @skipIfRocm
     def test_lazy_property_grad(self):
         x = torch.randn(1, requires_grad=True)
 
@@ -1129,6 +1133,7 @@ def test_poisson_sample(self):
 
     @unittest.skipIf(not TEST_CUDA, "CUDA not found")
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
+    @skipIfRocm
     def test_poisson_gpu_sample(self):
         set_rng_seed(1)
         for rate in [0.12, 0.9, 4.0]:
@@ -1536,6 +1541,7 @@ def test_normal_sample(self):
                                         scipy.stats.norm(loc=loc, scale=scale),
                                         'Normal(mean={}, std={})'.format(loc, scale))
 
+    @skipIfRocm
     def test_lowrank_multivariate_normal_shape(self):
         mean = torch.randn(5, 3, requires_grad=True)
         mean_no_batch = torch.randn(3, requires_grad=True)
@@ -1584,6 +1590,7 @@ def test_lowrank_multivariate_normal_shape(self):
                                  (mean_multi_batch, cov_factor_batched, cov_diag_batched))
 
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
+    @skipIfRocm
     def test_lowrank_multivariate_normal_log_prob(self):
         mean = torch.randn(3, requires_grad=True)
         cov_factor = torch.randn(3, 1, requires_grad=True)
@@ -1617,6 +1624,7 @@ def test_lowrank_multivariate_normal_log_prob(self):
         self.assertAlmostEqual(0.0, (batched_prob - unbatched_prob).abs().max(), places=3)
 
     @unittest.skipIf(not TEST_NUMPY, "NumPy not found")
+    @skipIfRocm
     def test_lowrank_multivariate_normal_sample(self):
         set_rng_seed(0)  # see Note [Randomized statistical tests]
         mean = torch.randn(5, requires_grad=True)
@@ -1629,6 +1637,7 @@ def test_lowrank_multivariate_normal_sample(self):
                                     'LowRankMultivariateNormal(loc={}, cov_factor={}, cov_diag={})'
                                     .format(mean, cov_factor, cov_diag), multivariate=True)
 
+    @skipIfRocm
     def test_lowrank_multivariate_normal_properties(self):
         loc = torch.randn(5)
         cov_factor = torch.randn(5, 2)
@@ -1643,6 +1652,7 @@ def test_lowrank_multivariate_normal_properties(self):
         self.assertEqual(m1.precision_matrix, m2.precision_matrix)
         self.assertEqual(m1.entropy(), m2.entropy())
 
+    @skipIfRocm
     def test_lowrank_multivariate_normal_moments(self):
         set_rng_seed(0)  # see Note [Randomized statistical tests]
         mean = torch.randn(5)
@@ -1655,6 +1665,7 @@ def test_lowrank_multivariate_normal_moments(self):
         empirical_var = samples.var(0)
         self.assertEqual(d.variance, empirical_var, prec=0.02)
 
+    @skipIfRocm
     def test_multivariate_normal_shape(self):
         mean = torch.randn(5, 3, requires_grad=True)
         mean_no_batch = torch.randn(3, requires_grad=True)
@@ -1702,6 +1713,7 @@ def test_multivariate_normal_shape(self):
         self._gradcheck_log_prob(MultivariateNormal, (mean_no_batch, None, None, scale_tril_batched))
 
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
+    @skipIfRocm
     def test_multivariate_normal_log_prob(self):
         mean = torch.randn(3, requires_grad=True)
         tmp = torch.randn(3, 10)
@@ -1739,6 +1751,7 @@ def test_multivariate_normal_log_prob(self):
         self.assertAlmostEqual(0.0, (batched_prob - unbatched_prob).abs().max(), places=3)
 
     @unittest.skipIf(not TEST_NUMPY, "NumPy not found")
+    @skipIfRocm
     def test_multivariate_normal_sample(self):
         set_rng_seed(0)  # see Note [Randomized statistical tests]
         mean = torch.randn(3, requires_grad=True)
@@ -1760,6 +1773,7 @@ def test_multivariate_normal_sample(self):
                                     'MultivariateNormal(loc={}, scale_tril={})'.format(mean, scale_tril),
                                     multivariate=True)
 
+    @skipIfRocm
     def test_multivariate_normal_properties(self):
         loc = torch.randn(5)
         scale_tril = transform_to(constraints.lower_cholesky)(torch.randn(5, 5))
@@ -1888,6 +1902,7 @@ def ref_log_prob(idx, x, log_prob):
 
     @unittest.skipIf(not TEST_CUDA, "CUDA not found")
     @unittest.skipIf(not TEST_NUMPY, "NumPy not found")
+    @skipIfRocm
     def test_gamma_gpu_shape(self):
         alpha = torch.tensor(torch.exp(torch.randn(2, 3).cuda()), requires_grad=True)
         beta = torch.tensor(torch.exp(torch.randn(2, 3).cuda()), requires_grad=True)
@@ -1918,6 +1933,7 @@ def test_gamma_sample(self):
 
     @unittest.skipIf(not TEST_CUDA, "CUDA not found")
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
+    @skipIfRocm
     def test_gamma_gpu_sample(self):
         set_rng_seed(0)
         for alpha, beta in product([0.1, 1.0, 5.0], [0.1, 1.0, 10.0]):
@@ -2150,6 +2166,7 @@ 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))
 
+    @skipIfRocm
     def test_independent_shape(self):
         for Dist, params in EXAMPLES:
             for i, param in enumerate(params):
@@ -2178,6 +2195,7 @@ def test_independent_shape(self):
                     except NotImplementedError:
                         pass
 
+    @skipIfRocm
     def test_cdf_icdf_inverse(self):
         # Tests the invertibility property on the distributions
         for Dist, params in EXAMPLES:
@@ -2197,6 +2215,7 @@ def test_cdf_icdf_inverse(self):
                     'icdf(cdf(x)) = {}'.format(actual),
                 ]))
 
+    @skipIfRocm
     def test_cdf_log_prob(self):
         # Tests if the differentiation of the CDF gives the PDF at a given value
         for Dist, params in EXAMPLES:
@@ -2588,6 +2607,7 @@ def tearDown(self):
         super(TestCase, self).tearDown()
         Distribution.set_default_validate_args(False)
 
+    @skipIfRocm
     def test_entropy_shape(self):
         for Dist, params in EXAMPLES:
             for i, param in enumerate(params):
@@ -3141,6 +3161,7 @@ def test_kl_monte_carlo(self):
 
     # Multivariate normal has a separate Monte Carlo based test due to the requirement of random generation of
     # positive (semi) definite matrices. n is set to 5, but can be increased during testing.
+    @skipIfRocm
     def test_kl_multivariate_normal(self):
         set_rng_seed(0)  # see Note [Randomized statistical tests]
         n = 5  # Number of tests for multivariate_normal
@@ -3166,6 +3187,7 @@ def test_kl_multivariate_normal(self):
                 'Actual (analytic): {}'.format(actual),
             ]))
 
+    @skipIfRocm
     def test_kl_multivariate_normal_batched(self):
         b = 7  # Number of batches
         loc = [torch.randn(b, 3) for _ in range(0, 2)]
@@ -3177,6 +3199,7 @@ def test_kl_multivariate_normal_batched(self):
                                   MultivariateNormal(loc[1], scale_tril=scale_tril[1]))
         self.assertEqual(expected_kl, actual_kl)
 
+    @skipIfRocm
     def test_kl_multivariate_normal_batched_broadcasted(self):
         b = 7  # Number of batches
         loc = [torch.randn(b, 3) for _ in range(0, 2)]
@@ -3189,6 +3212,7 @@ def test_kl_multivariate_normal_batched_broadcasted(self):
                                   MultivariateNormal(loc[1], scale_tril=scale_tril[1]))
         self.assertEqual(expected_kl, actual_kl)
 
+    @skipIfRocm
     def test_kl_lowrank_multivariate_normal(self):
         set_rng_seed(0)  # see Note [Randomized statistical tests]
         n = 5  # Number of tests for lowrank_multivariate_normal
@@ -3229,6 +3253,7 @@ def test_kl_lowrank_multivariate_normal(self):
                 'Actual (analytic): {}'.format(actual_full_lowrank),
             ]))
 
+    @skipIfRocm
     def test_kl_lowrank_multivariate_normal_batched(self):
         b = 7  # Number of batches
         loc = [torch.randn(b, 3) for _ in range(0, 2)]
@@ -3264,6 +3289,7 @@ def test_kl_edgecases(self):
         self.assertEqual(kl_divergence(Bernoulli(1), Bernoulli(1)), 0)
         self.assertEqual(kl_divergence(Categorical(torch.tensor([0., 1.])), Categorical(torch.tensor([0., 1.]))), 0)
 
+    @skipIfRocm
     def test_kl_shape(self):
         for Dist, params in EXAMPLES:
             for i, param in enumerate(params):
@@ -3279,6 +3305,7 @@ def test_kl_shape(self):
                     'Actual {}'.format(kl.shape),
                 ]))
 
+    @skipIfRocm
     def test_entropy_monte_carlo(self):
         set_rng_seed(0)  # see Note [Randomized statistical tests]
         for Dist, params in EXAMPLES:
@@ -3322,6 +3349,7 @@ def test_entropy_exponential_family(self):
 
 
 class TestConstraints(TestCase):
+    @skipIfRocm
     def test_params_contains(self):
         for Dist, params in EXAMPLES:
             for i, param in enumerate(params):
@@ -3345,6 +3373,7 @@ def test_params_contains(self):
                         Dist.__name__, i + 1, len(params), name, value)
                     self.assertTrue(constraint.check(value).all(), msg=message)
 
+    @skipIfRocm
     def test_support_contains(self):
         for Dist, params in EXAMPLES:
             self.assertIsInstance(Dist.support, Constraint)
@@ -3624,6 +3653,7 @@ def setUp(self):
             )
         ]
 
+    @skipIfRocm
     def test_mean(self):
         for pytorch_dist, scipy_dist in self.distribution_pairs:
             if isinstance(pytorch_dist, (Cauchy, HalfCauchy)):
@@ -3634,6 +3664,7 @@ def test_mean(self):
             else:
                 self.assertEqual(pytorch_dist.mean, scipy_dist.mean(), allow_inf=True, message=pytorch_dist)
 
+    @skipIfRocm
     def test_variance_stddev(self):
         for pytorch_dist, scipy_dist in self.distribution_pairs:
             if isinstance(pytorch_dist, (Cauchy, HalfCauchy)):
@@ -3649,6 +3680,7 @@ def test_variance_stddev(self):
                 self.assertEqual(pytorch_dist.variance, scipy_dist.var(), allow_inf=True, message=pytorch_dist)
                 self.assertEqual(pytorch_dist.stddev, scipy_dist.var() ** 0.5, message=pytorch_dist)
 
+    @skipIfRocm
     def test_cdf(self):
         for pytorch_dist, scipy_dist in self.distribution_pairs:
             samples = pytorch_dist.sample((5,))
@@ -3658,6 +3690,7 @@ def test_cdf(self):
                 continue
             self.assertEqual(cdf, scipy_dist.cdf(samples), message=pytorch_dist)
 
+    @skipIfRocm
     def test_icdf(self):
         for pytorch_dist, scipy_dist in self.distribution_pairs:
             samples = torch.rand((5,) + pytorch_dist.batch_shape)
@@ -3963,6 +3996,7 @@ def test_biject_to(self):
             self.assertEqual(j.shape, x.shape[:x.dim() - t.event_dim])
 
     @unittest.skipIf(not TEST_CUDA, "CUDA not found")
+    @skipIfRocm
     def test_biject_to_cuda(self):
         for constraint in self.get_constraints(is_cuda=True):
             try:
@@ -3995,6 +4029,7 @@ def test_transform_to(self):
             self.assertEqual(y, y2, message="Error in transform_to({}) pseudoinverse".format(constraint))
 
     @unittest.skipIf(not TEST_CUDA, "CUDA not found")
+    @skipIfRocm
     def test_transform_to_cuda(self):
         for constraint in self.get_constraints(is_cuda=True):
             t = transform_to(constraint)
@@ -4012,6 +4047,7 @@ def setUp(self):
         super(TestCase, self).setUp()
         Distribution.set_default_validate_args(True)
 
+    @skipIfRocm
     def test_valid(self):
         for Dist, params in EXAMPLES:
             for i, param in enumerate(params):