Fix variance shape bug in Riemann posterior (#2939)

Summary:
Pull Request resolved: #2939

Currently, BoundedRiemannPosterior.mean returns a Size([b, 1]) tensor, but BoundedRiemannPosterior.mean_of_square returns a Size([b]) tensor. As a result, when these are combined to get the variance, we end up with a Size([b, b]) variance instead of the correct Size([b, 1]). This fixes the issue.

Reviewed By: SamuelGabriel

Differential Revision: D78911131

fbshipit-source-id: 5ba69604b3c198ddce68b70a4af806e64fdbc0b1

Author: bletham

botorch_community/posteriors/riemann.py

@@ -107,7 +107,7 @@ def mean(self):
         r"""The mean of the posterior distribution."""
         bucket_widths = self.borders[1:] - self.borders[:-1]
         bucket_means = self.borders[:-1] + bucket_widths / 2
-        return (bucket_means * (self.probabilities)).sum(-1, keepdim=True)
+        return (self.probabilities @ bucket_means).unsqueeze(-1)
 
     @property
     def mean_of_square(self) -> torch.Tensor:
@@ -119,7 +119,7 @@ def mean_of_square(self) -> torch.Tensor:
             + right_borders.square()
             + left_borders * right_borders
         ) / 3.0
-        return self.probabilities @ bucket_mean_of_square
+        return (self.probabilities @ bucket_mean_of_square).unsqueeze(-1)
 
     @property
     def variance(self) -> torch.Tensor:

test_community/posteriors/test_riemann.py

@@ -148,6 +148,13 @@ def test_variance(self):
             computed_variance = posterior.variance
             self.assertLess((computed_variance - true_variance).abs().item(), 0.05)
 
+            # Check with batch dimension
+            probabilities = torch.rand(2, n_buckets, **tkwargs)
+            probabilities = probabilities / probabilities.sum(-1, keepdim=True)
+            posterior = BoundedRiemannPosterior(borders, probabilities)
+            self.assertEqual(posterior.variance.shape, torch.Size([2, 1]))
+            self.assertEqual(posterior.mean.shape, torch.Size([2, 1]))
+
     def test_confidence_region(self):
         torch.manual_seed(13)
         for dtype in (torch.float, torch.double):