Merge pull request #2621 from saitcakmak/mtmvn

Support broadcasting batch shapes in MTMVN.from_independent_mvns

Author: saitcakmak

gpytorch/distributions/multitask_multivariate_normal.py

@@ -153,7 +153,8 @@ def from_independent_mvns(cls, mvns):
         if any(isinstance(mvn, MultitaskMultivariateNormal) for mvn in mvns):
             raise ValueError("Cannot accept MultitaskMultivariateNormals")
         if not all(m.batch_shape == mvns[0].batch_shape for m in mvns[1:]):
-            raise ValueError("All MultivariateNormals must have the same batch shape")
+            batch_shape = torch.broadcast_shapes(*(m.batch_shape for m in mvns))
+            mvns = [mvn.expand(batch_shape) for mvn in mvns]
         if not all(m.event_shape == mvns[0].event_shape for m in mvns[1:]):
             raise ValueError("All MultivariateNormals must have the same event shape")
         mean = torch.stack([mvn.mean for mvn in mvns], -1)

test/distributions/test_multitask_multivariate_normal.py

@@ -278,6 +278,14 @@ def test_from_independent_mvns(self, cuda=False):
             self.assertEqual(list(mvn.mean.shape), expected_mean_shape)
             self.assertEqual(list(mvn.covariance_matrix.shape), expected_covar_shape)
 
+            # Test mixed batch mode mvns
+            # Second MVN is batched, so the first one will be expanded to match.
+            mvns[1] = mvns[1].expand(torch.Size([3]))
+            expected_mvn = mvn.expand(torch.Size([3]))
+            mvn = MultitaskMultivariateNormal.from_independent_mvns(mvns=mvns)
+            self.assertTrue(torch.equal(mvn.mean, expected_mvn.mean))
+            self.assertTrue(torch.equal(mvn.covariance_matrix, expected_mvn.covariance_matrix))
+
             # Test batch mode mvns
             b = 3
             mvns = [