add dim validation

gpytorch/distributions/multivariate_normal.py

@@ -170,8 +170,13 @@ def unsqueeze(self, dim: int) -> MultivariateNormal:
         If `dim = -1`, then the returned MultivariateNormal will have
         `batch_shape = torch.Size([2, 3, 1])`.
         """
-        # If dim is negative, get the positive equivalent.
+        if dim > len(self.batch_shape) or dim < -len(self.batch_shape) - 1:
+            raise IndexError(
+                "Dimension out of range (expected to be in range of "
+                f"[{-len(self.batch_shape) - 1}, {len(self.batch_shape)}], but got {dim})."
+            )
         if dim < 0:
+            # If dim is negative, get the positive equivalent.
             dim = len(self.batch_shape) + dim + 1
 
         new_loc = self.loc.unsqueeze(dim)

test/distributions/test_multivariate_normal.py

@@ -370,8 +370,20 @@ def test_multivariate_normal_unsqueeze(self, cuda=False):
                 self.assertEqual(new.batch_shape, expected_batch)
                 self.assertEqual(new.event_shape, mvn.event_shape)
                 self.assertTrue(torch.equal(new.mean, mean.unsqueeze(positive_dim)))
+                self.assertEqual(new.mean.shape, expected_batch + torch.Size([3]))
                 self.assertTrue(torch.allclose(new.covariance_matrix, covmat.unsqueeze(positive_dim)))
+                self.assertEqual(new.covariance_matrix.shape, expected_batch + torch.Size([3, 3]))
                 self.assertTrue(torch.allclose(new.scale_tril, mvn.scale_tril.unsqueeze(positive_dim)))
+                self.assertEqual(new.scale_tril.shape, expected_batch + torch.Size([3, 3]))
+
+        # Check for dim validation.
+        with self.assertRaisesRegex(IndexError, "Dimension out of range"):
+            mvn.unsqueeze(3)
+        with self.assertRaisesRegex(IndexError, "Dimension out of range"):
+            mvn.unsqueeze(-4)
+        # Should not raise error up to 2 or -3.
+        mvn.unsqueeze(2)
+        mvn.unsqueeze(-3)
 
 
 if __name__ == "__main__":