expand covar matrix

cornellius-gp · Jan 21, 2025 · b37a85b · b37a85b
1 parent 9eaecdc
commit b37a85b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/gpytorch/distributions/multivariate_normal.py b/gpytorch/distributions/multivariate_normal.py
@@ -158,7 +158,7 @@ def expand(self, batch_size: torch.Size) -> MultivariateNormal:
             )
             super(MultivariateNormal, new).__init__(loc=new_loc, scale_tril=new_scale_tril)
             # Set the covar matrix, since it is always available for GPyTorch MVN.
-            new.covariance_matrix = self.covariance_matrix
+            new.covariance_matrix = self.covariance_matrix.expand(batch_size + self.covariance_matrix.shape[-2:])
         return new
 
     def get_base_samples(self, sample_shape: torch.Size = torch.Size()) -> Tensor:

diff --git a/test/distributions/test_multivariate_normal.py b/test/distributions/test_multivariate_normal.py
@@ -343,8 +343,11 @@ def test_multivariate_normal_expand(self, cuda=False):
             self.assertEqual(expanded.batch_shape, torch.Size([2]))
             self.assertEqual(expanded.event_shape, mvn.event_shape)
             self.assertTrue(torch.equal(expanded.mean, mean.expand(2, -1)))
+            self.assertEqual(expanded.mean.shape, torch.Size([2, 3]))
             self.assertTrue(torch.allclose(expanded.covariance_matrix, covmat.expand(2, -1, -1)))
+            self.assertEqual(expanded.covariance_matrix.shape, torch.Size([2, 3, 3]))
             self.assertTrue(torch.allclose(expanded.scale_tril, mvn.scale_tril.expand(2, -1, -1)))
+            self.assertEqual(expanded.scale_tril.shape, torch.Size([2, 3, 3]))
 
 
 if __name__ == "__main__":