Back out "Back out "[botorch] Replace Pyro with NumPyro for fully Bayesian NUTS inference (96% reduction in fit time)"" (#5162)

Summary:
X-link: meta-pytorch/botorch#3269

Pull Request resolved: #5162

Reviewed By: saitcakmak

Differential Revision: D99688719

fbshipit-source-id: d78f3f75dd0503b61e1cb7f48e64484979c3ee48

Author: sdaulton

ax/utils/sensitivity/tests/test_sensitivity.py

@@ -201,37 +201,23 @@ def test_DgsmGpSampling(self) -> None:
     def _test_sobol_gp_mean(
         self,
         sensitivity: SobolSensitivityGPMean,
-        expected_first_order: Tensor,
-        expected_total_order: Tensor,
-        expected_second_order: Tensor | None = None,
+        expected_first_order_shape: torch.Size,
+        expected_total_order_shape: torch.Size,
+        expected_second_order_shape: torch.Size | None = None,
     ) -> None:
-        """
-        Check that outputs are as expected. The `assertAllClose` checks are
-        characterization tests rather than correctness tests; they check that
-        behavior is the same as it was in the past, not that it is
-        quantitatively correct. Innocuous changes such as changing a random seed
-        could potentially break these tests, and it may become necessary to
-        delete them.
-        """
-        atol = 4e-3
-        rtol = 2e-3
+        """Check that outputs have the correct type and shape."""
         first_order = sensitivity.first_order_indices()
         self.assertIsInstance(first_order, Tensor)
-        self.assertAllClose(first_order, expected_first_order, atol=atol, rtol=rtol)
-        self.assertEqual(first_order.shape, expected_first_order.shape)
+        self.assertEqual(first_order.shape, expected_first_order_shape)
 
         total_order = sensitivity.total_order_indices()
         self.assertIsInstance(total_order, Tensor)
-        self.assertAllClose(total_order, expected_total_order, atol=atol, rtol=rtol)
-        self.assertEqual(total_order.shape, expected_total_order.shape)
+        self.assertEqual(total_order.shape, expected_total_order_shape)
 
-        if expected_second_order is not None:
+        if expected_second_order_shape is not None:
             second_order = sensitivity.second_order_indices()
             self.assertIsInstance(second_order, Tensor)
-            self.assertAllClose(
-                second_order, expected_second_order, atol=atol, rtol=rtol
-            )
-            self.assertEqual(second_order.shape, expected_second_order.shape)
+            self.assertEqual(second_order.shape, expected_second_order_shape)
 
     def test_SobolGPMean(self) -> None:
         bounds = torch.tensor([(0.0, 1.0) for _ in range(2)]).t()
@@ -240,9 +226,9 @@ def test_SobolGPMean(self) -> None:
         )
         self._test_sobol_gp_mean(
             sensitivity=sensitivity_mean,
-            expected_first_order=torch.tensor([1.1547, -0.4024], dtype=torch.float64),
-            expected_total_order=torch.tensor([0.4299, 0.4894], dtype=torch.float64),
-            expected_second_order=torch.tensor([-1.4845], dtype=torch.float64),
+            expected_first_order_shape=torch.Size([2]),
+            expected_total_order_shape=torch.Size([2]),
+            expected_second_order_shape=torch.Size([1]),
         )
 
     def test_SobolGPMean_SAASBO(self) -> None:
@@ -252,9 +238,9 @@ def test_SobolGPMean_SAASBO(self) -> None:
         )
         self._test_sobol_gp_mean(
             sensitivity=sensitivity_mean_saas,
-            expected_first_order=torch.tensor([0.5752, 0.5143], dtype=torch.double),
-            expected_total_order=torch.tensor([0.9897, 0.0979], dtype=torch.float64),
-            expected_second_order=torch.tensor([0.8332], dtype=torch.double),
+            expected_first_order_shape=torch.Size([2]),
+            expected_total_order_shape=torch.Size([2]),
+            expected_second_order_shape=torch.Size([1]),
         )
 
         sensitivity_mean_bootstrap = SobolSensitivityGPMean(
@@ -267,17 +253,9 @@ def test_SobolGPMean_SAASBO(self) -> None:
         )
         self._test_sobol_gp_mean(
             sensitivity=sensitivity_mean_bootstrap,
-            expected_first_order=torch.tensor(
-                [[0.6327, 10.0889, 1.0044], [0.2089, 0.7322, 0.2706]],
-                dtype=torch.float64,
-            ),
-            expected_total_order=torch.tensor(
-                [[0.8013, 0.1824, 0.1351], [0.2203, 0.0304, 0.0551]],
-                dtype=torch.float64,
-            ),
-            expected_second_order=torch.tensor(
-                [[0.7978, 22.6598, 1.5053]], dtype=torch.float64
-            ),
+            expected_first_order_shape=torch.Size([2, 3]),
+            expected_total_order_shape=torch.Size([2, 3]),
+            expected_second_order_shape=torch.Size([1, 3]),
         )
 
         sensitivity_mean_bootstrap = SobolSensitivityGPMean(
@@ -290,26 +268,18 @@ def test_SobolGPMean_SAASBO(self) -> None:
         )
         self._test_sobol_gp_mean(
             sensitivity=sensitivity_mean_bootstrap,
-            expected_first_order=torch.tensor(
-                [[3.4512, 32.4428, 1.8012], [0.2069, 121.8610, 3.4909]],
-                dtype=torch.float64,
-            ),
-            expected_total_order=torch.tensor(
-                [[0.4288, 0.0903, 0.0950], [0.7923, 0.2218, 0.1489]],
-                dtype=torch.float64,
-            ),
-            expected_second_order=torch.tensor(
-                [[-6.3790, 397.4363, 6.3043]], dtype=torch.float64
-            ),
+            expected_first_order_shape=torch.Size([2, 3]),
+            expected_total_order_shape=torch.Size([2, 3]),
+            expected_second_order_shape=torch.Size([1, 3]),
         )
 
         sensitivity_mean = SobolSensitivityGPMean(
             self.model, num_mc_samples=10, bounds=bounds, second_order=False
         )
         self._test_sobol_gp_mean(
             sensitivity=sensitivity_mean,
-            expected_first_order=torch.tensor([0.9566, -0.4183], dtype=torch.float64),
-            expected_total_order=torch.tensor([0.3440, 0.3685], dtype=torch.float64),
+            expected_first_order_shape=torch.Size([2]),
+            expected_total_order_shape=torch.Size([2]),
         )
 
         with self.assertRaisesRegex(ValueError, "Second order indices"):

ax/utils/testing/tests/test_mock.py

@@ -8,6 +8,7 @@
 
 from unittest.mock import patch
 
+import jax.numpy as jnp
 import torch
 from ax.adapter.registry import Generators
 from ax.adapter.transforms.choice_encode import OrderedChoiceToIntegerRange
@@ -22,7 +23,6 @@
 from botorch.generation.gen import gen_candidates_scipy
 from botorch.optim.optimize_mixed import generate_starting_points
 from botorch.utils.testing import MockAcquisitionFunction, skip_if_import_error
-from pyro.infer import MCMC
 
 
 class TestMock(TestCase):
@@ -45,13 +45,25 @@ def test_botorch_mocks(self) -> None:
 
     def test_fully_bayesian_mocks(self) -> None:
         experiment = get_branin_experiment(with_completed_batch=True)
-        with patch("botorch.fit.MCMC", wraps=MCMC) as mock_mcmc:
+        num_mcmc_samples = 16
+        dim = len(experiment.search_space.parameters)
+        # Mock MCMC to return proper JAX arrays that postprocess_mcmc_samples
+        # can handle (it performs JAX operations on them).
+        mock_samples = {
+            "mean": jnp.ones(num_mcmc_samples),
+            "noise": jnp.ones(num_mcmc_samples),
+            "outputscale": jnp.ones(num_mcmc_samples),
+            "kernel_tausq": jnp.ones(num_mcmc_samples),
+            "_kernel_inv_length_sq": jnp.ones((num_mcmc_samples, dim)),
+        }
+        with patch("botorch.fit.MCMC") as mock_mcmc:
+            mock_mcmc.return_value.get_samples.return_value = mock_samples
             with mock_botorch_optimize_context_manager():
                 Generators.SAASBO(experiment=experiment, data=experiment.lookup_data())
         mock_mcmc.assert_called_once()
         kwargs = mock_mcmc.call_args.kwargs
         self.assertEqual(kwargs["num_samples"], 16)
-        self.assertEqual(kwargs["warmup_steps"], 0)
+        self.assertEqual(kwargs["num_warmup"], 0)
 
     def test_mixed_optimizer_mocks(self) -> None:
         experiment = get_branin_experiment(