Fix broadcasting bug in StudentT (#12148)

fritzo · facebook-github-bot · commit 0aff3cc55942 · 2018-09-28T09:57:02.000-07:00
Summary: This fixes a broadcasting error with the `StudentT` distribution - [x] added a regression test - [x] strengthened parameter broadcasting tests Pull Request resolved: #12148 Differential Revision: D10099226 Pulled By: soumith fbshipit-source-id: 0c5eb14180d158f8fff28ceb9e7cd3471c2bb803
diff --git a/test/test_distributions.py b/test/test_distributions.py
@@ -2378,12 +2378,20 @@ def test_valid_parameter_broadcasting(self):
              (1, 2)),
             (StudentT(df=torch.tensor([1.]), scale=torch.tensor([[1.]])),
              (1, 1)),
+            (StudentT(df=1., loc=torch.zeros(5, 1), scale=torch.ones(3)),
+             (5, 3)),
         ]
 
         for dist, expected_size in valid_examples:
-            dist_sample_size = dist.sample().size()
-            self.assertEqual(dist_sample_size, expected_size,
-                             'actual size: {} != expected size: {}'.format(dist_sample_size, expected_size))
+            actual_size = dist.sample().size()
+            self.assertEqual(actual_size, expected_size,
+                             '{} actual size: {} != expected size: {}'.format(dist, actual_size, expected_size))
+
+            sample_shape = torch.Size((2,))
+            expected_size = sample_shape + expected_size
+            actual_size = dist.sample(sample_shape).size()
+            self.assertEqual(actual_size, expected_size,
+                             '{} actual size: {} != expected size: {}'.format(dist, actual_size, expected_size))
 
     def test_invalid_parameter_broadcasting(self):
         # invalid broadcasting cases; should throw error
diff --git a/torch/distributions/studentT.py b/torch/distributions/studentT.py
@@ -41,8 +41,8 @@ def variance(self):
 
     def __init__(self, df, loc=0., scale=1., validate_args=None):
         self.df, self.loc, self.scale = broadcast_all(df, loc, scale)
-        self._chi2 = Chi2(df)
-        batch_shape = torch.Size() if isinstance(df, Number) else self.df.size()
+        self._chi2 = Chi2(self.df)
+        batch_shape = self.df.size()
         super(StudentT, self).__init__(batch_shape, validate_args=validate_args)
 
     def expand(self, batch_shape, _instance=None):
