Weighted & batched EnsemblePosterior Ax support (facebook#4201)

Summary:
Adds more predict support for `EnsemblePosterior`, and makes it more coherent with the predict support that is already in place for `GaussianMixturePosterior` by using the newly added (previous diff) `mixture` attributes for `EnsemblePosterior`.

In the case where an `EnsemblePosterior` has batch dimensions in addition to the ensemble dimension, we compute the mean and variance over both batch and ensemble dimensions, instead of just the latter. Moreover, this allows us to properly account for non-uniform weights across batch and ensemble dimensions, which is needed for sampling-based benchmarking in ensembled Fully Bayesian models.

Removed tests that involved previous mixture moment calculations, as these are now in BoTorch.

Pull Request resolved: facebook#4201

Reviewed By: saitcakmak

Differential Revision: D80972578

fbshipit-source-id: 179799158979078958b9e9ba958d6ef58d996873

Author: Carl Hvarfner

ax/generators/torch/tests/test_utils.py

@@ -736,85 +736,21 @@ def test_model_config(self) -> None:
 
     def test_predict_from_model_ensemble_posterior(self) -> None:
         """Test predict_from_model with EnsemblePosterior support."""
-        # Create test data
         X = torch.rand(2, 3)
 
-        # Create a mock EnsemblePosterior with ndim > 2
+        # Create a mock EnsemblePosterior
         mock_posterior = Mock(spec=EnsemblePosterior)
-
-        # Set up posterior values with shape (num_models, batch_shape, output_shape)
-        # This simulates an ensemble of 5 models with 2 test points and 2 outputs
-        posterior_values = torch.rand(5, 2, 2)  # (5 models, 2 points, 2 outputs)
-        mock_posterior.values = posterior_values
+        mock_posterior.mixture_mean = torch.rand(2, 2)
+        mock_posterior.mixture_variance = torch.rand(2, 2)
 
         # Create a mock model
         mock_model = Mock()
         mock_model.posterior.return_value = mock_posterior
 
-        # Test with use_posterior_predictive=False
+        # Test prediction
         mean, cov = predict_from_model(mock_model, X, use_posterior_predictive=False)
 
-        # Verify the model.posterior was called correctly
-        mock_model.posterior.assert_called_once()
-
-        # Verify output shapes
+        # Verify shapes
         self.assertEqual(mean.shape, (2, 2))  # (n_points, n_outputs)
         self.assertEqual(cov.shape, (2, 2, 2))  # (n_points, n_outputs, n_outputs)
-
-        # Verify mean calculation (should be mean over ensemble dimension)
-        expected_mean = posterior_values.mean(dim=0)  # Average over first dimension
-        self.assertTrue(torch.allclose(mean, expected_mean))
-
-        # Verify variance calculation (should be variance over ensemble dimension)
-        expected_var = posterior_values.var(dim=0)
-        # Check that the diagonal of the covariance matches expected variance
-        self.assertTrue(
-            torch.allclose(torch.diagonal(cov, dim1=-2, dim2=-1), expected_var)
-        )
-
-        # Test with use_posterior_predictive=True
-        mock_model.reset_mock()
-        predict_from_model(mock_model, X, use_posterior_predictive=True)
-        mock_model.posterior.assert_called_once()
-
-        mock_posterior2 = Mock(spec=EnsemblePosterior)
-        # Shape: (num_models, batch1, batch2, output_shape) - ndim = 4
-        posterior_values_5d = torch.rand(2, 3, 4, 2, 2)
-        mock_posterior2.values = posterior_values_5d
-        mock_model2 = Mock()
-        mock_model2.posterior.return_value = mock_posterior2
-
-        X2 = torch.rand(4, 3)
-        mean2, cov2 = predict_from_model(
-            mock_model2, X2, use_posterior_predictive=False
-        )
-
-        # Should average over first two dimensions (all except last 2)
-        expected_mean_5d = posterior_values_5d.mean(dim=(0, 1, 2))
-        expected_var_5d = posterior_values_5d.var(dim=(0, 1, 2))
-
-        self.assertTrue(torch.allclose(mean2, expected_mean_5d))
-        self.assertTrue(
-            torch.allclose(torch.diagonal(cov2, dim1=-2, dim2=-1), expected_var_5d)
-        )
-
-        # Test case where ensemble size is 1 or non-existant
-        # (variance should be zero, not NaN)
-        posterior_values_singles = [
-            torch.rand(1, 2, 2),
-            torch.rand(1, 1, 2, 2),
-        ]  # Single ensemble model
-        mock_model.reset_mock()
-        mock_posterior_single = Mock(spec=EnsemblePosterior)
-        mock_model.posterior.return_value = mock_posterior_single
-        for i, posterior_values_single in enumerate(posterior_values_singles):
-            with self.subTest(i=i, shape=posterior_values_single.shape):
-                mock_posterior_single.values = posterior_values_single
-                mean_single, cov_single = predict_from_model(
-                    mock_model, X, use_posterior_predictive=False
-                )
-                # Variance should be zero (not NaN) when ensemble size is 1
-                var_single = torch.diagonal(cov_single, dim1=-2, dim2=-1)
-                self.assertTrue(
-                    torch.allclose(var_single, torch.zeros_like(var_single))
-                )
+        self.assertTrue(torch.all(cov >= 0))  # Ensure covariance is positive

ax/generators/torch/utils.py

@@ -535,17 +535,11 @@ def predict_from_model(
             if isinstance(posterior, GaussianMixturePosterior):
                 mean = posterior.mixture_mean.cpu().detach()
                 var = posterior.mixture_variance.cpu().detach().clamp_min(0)
-            elif isinstance(posterior, EnsemblePosterior) and posterior.values.ndim > 2:
-                # Compute dimensions to average over (all except last 2)
-                # Not using the build-in EnsemblePosterior.variance() since that
-                # does not allow us to compute variance over _all_ batch dimensions
-                avg_dims = tuple(range(posterior.values.ndim - 2))
-                mean = posterior.values.mean(dim=avg_dims).cpu().detach()
-                var = posterior.values.var(dim=avg_dims).cpu().detach()
-
-                # Replace NaN values with zero (occurs when ensemble size is 1)
-                if posterior.values[..., 0, 0].numel() == 1:
-                    var = torch.zeros_like(var)
+            elif isinstance(posterior, EnsemblePosterior):
+                # Always use mixture_mean and mixture_variance for ensemble
+                # predictions - provides prediction from mixture, not just average
+                mean = posterior.mixture_mean.cpu().detach()
+                var = posterior.mixture_variance.cpu().detach().clamp_min(0)
             else:
                 try:
                     mean = posterior.mean.cpu().detach()  # type: ignore