Back out "Replace Pyro with NumPyro for fully Bayesian NUTS inference (96% reduction in fit time)" (facebook#5136)

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

Pull Request resolved: facebook#5136

Reviewed By: Balandat

Differential Revision: D99446023

fbshipit-source-id: 76efd2e595d48589ee3252efe407e31521044468

Author: sdaulton

ax/utils/sensitivity/tests/test_sensitivity.py

@@ -201,23 +201,37 @@ def test_DgsmGpSampling(self) -> None:
     def _test_sobol_gp_mean(
         self,
         sensitivity: SobolSensitivityGPMean,
-        expected_first_order_shape: torch.Size,
-        expected_total_order_shape: torch.Size,
-        expected_second_order_shape: torch.Size | None = None,
+        expected_first_order: Tensor,
+        expected_total_order: Tensor,
+        expected_second_order: Tensor | None = None,
     ) -> None:
-        """Check that outputs have the correct type and shape."""
+        """
+        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
         first_order = sensitivity.first_order_indices()
         self.assertIsInstance(first_order, Tensor)
-        self.assertEqual(first_order.shape, expected_first_order_shape)
+        self.assertAllClose(first_order, expected_first_order, atol=atol, rtol=rtol)
+        self.assertEqual(first_order.shape, expected_first_order.shape)
 
         total_order = sensitivity.total_order_indices()
         self.assertIsInstance(total_order, Tensor)
-        self.assertEqual(total_order.shape, expected_total_order_shape)
+        self.assertAllClose(total_order, expected_total_order, atol=atol, rtol=rtol)
+        self.assertEqual(total_order.shape, expected_total_order.shape)
 
-        if expected_second_order_shape is not None:
+        if expected_second_order is not None:
             second_order = sensitivity.second_order_indices()
             self.assertIsInstance(second_order, Tensor)
-            self.assertEqual(second_order.shape, expected_second_order_shape)
+            self.assertAllClose(
+                second_order, expected_second_order, atol=atol, rtol=rtol
+            )
+            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()
@@ -226,9 +240,9 @@ def test_SobolGPMean(self) -> None:
         )
         self._test_sobol_gp_mean(
             sensitivity=sensitivity_mean,
-            expected_first_order_shape=torch.Size([2]),
-            expected_total_order_shape=torch.Size([2]),
-            expected_second_order_shape=torch.Size([1]),
+            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),
         )
 
     def test_SobolGPMean_SAASBO(self) -> None:
@@ -238,9 +252,9 @@ def test_SobolGPMean_SAASBO(self) -> None:
         )
         self._test_sobol_gp_mean(
             sensitivity=sensitivity_mean_saas,
-            expected_first_order_shape=torch.Size([2]),
-            expected_total_order_shape=torch.Size([2]),
-            expected_second_order_shape=torch.Size([1]),
+            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),
         )
 
         sensitivity_mean_bootstrap = SobolSensitivityGPMean(
@@ -253,9 +267,17 @@ def test_SobolGPMean_SAASBO(self) -> None:
         )
         self._test_sobol_gp_mean(
             sensitivity=sensitivity_mean_bootstrap,
-            expected_first_order_shape=torch.Size([2, 3]),
-            expected_total_order_shape=torch.Size([2, 3]),
-            expected_second_order_shape=torch.Size([1, 3]),
+            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
+            ),
         )
 
         sensitivity_mean_bootstrap = SobolSensitivityGPMean(
@@ -268,18 +290,26 @@ def test_SobolGPMean_SAASBO(self) -> None:
         )
         self._test_sobol_gp_mean(
             sensitivity=sensitivity_mean_bootstrap,
-            expected_first_order_shape=torch.Size([2, 3]),
-            expected_total_order_shape=torch.Size([2, 3]),
-            expected_second_order_shape=torch.Size([1, 3]),
+            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
+            ),
         )
 
         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_shape=torch.Size([2]),
-            expected_total_order_shape=torch.Size([2]),
+            expected_first_order=torch.tensor([0.9566, -0.4183], dtype=torch.float64),
+            expected_total_order=torch.tensor([0.3440, 0.3685], dtype=torch.float64),
         )
 
         with self.assertRaisesRegex(ValueError, "Second order indices"):

ax/utils/testing/tests/test_mock.py

@@ -8,7 +8,6 @@
 
 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
@@ -23,6 +22,7 @@
 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,25 +45,13 @@ def test_botorch_mocks(self) -> None:
 
     def test_fully_bayesian_mocks(self) -> None:
         experiment = get_branin_experiment(with_completed_batch=True)
-        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 patch("botorch.fit.MCMC", wraps=MCMC) as mock_mcmc:
             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["num_warmup"], 0)
+        self.assertEqual(kwargs["warmup_steps"], 0)
 
     def test_mixed_optimizer_mocks(self) -> None:
         experiment = get_branin_experiment(