Back out "Replace Pyro with NumPyro for fully Bayesian NUTS inference (96% reduction in fit time)" (meta-pytorch#3263)

Summary:
Pull Request resolved: meta-pytorch#3263

X-link: facebook/Ax#5136

Reviewed By: Balandat

Differential Revision: D99446023

fbshipit-source-id: 76efd2e595d48589ee3252efe407e31521044468

Author: sdaulton

botorch/fit.py

@@ -15,7 +15,6 @@
 from typing import Any
 from warnings import catch_warnings, simplefilter, warn_explicit, WarningMessage
 
-import jax
 from botorch.exceptions.errors import ModelFittingError, UnsupportedError
 from botorch.exceptions.warnings import OptimizationWarning
 from botorch.logging import logger
@@ -44,7 +43,7 @@
 from gpytorch.mlls.marginal_log_likelihood import MarginalLogLikelihood
 from gpytorch.mlls.sum_marginal_log_likelihood import SumMarginalLogLikelihood
 from linear_operator.utils.errors import NotPSDError
-from numpyro.infer import MCMC, NUTS
+from pyro.infer.mcmc import MCMC, NUTS
 from torch import device, Tensor
 from torch.nn import Parameter
 from torch.utils.data import DataLoader
@@ -343,11 +342,9 @@ def fit_fully_bayesian_model_nuts(
     thinning: int = 16,
     disable_progbar: bool = False,
     jit_compile: bool = False,
-    seed: int = 0,
 ) -> None:
     r"""Fit a fully Bayesian model using the No-U-Turn-Sampler (NUTS)
 
-    Uses NumPyro's NUTS implementation (backed by JAX) for MCMC inference.
 
     Args:
         model: Fully Bayesian GP to be fitted.
@@ -360,8 +357,7 @@ def fit_fully_bayesian_model_nuts(
             bar and diagnostics during MCMC.
         jit_compile: Whether to use jit. Using jit may be ~2X faster (rough estimate),
             but it will also increase the memory usage and sometimes result in runtime
-            errors.
-        seed: Random seed for JAX PRNG.
+            errors, e.g., https://github.com/pyro-ppl/pyro/issues/3136.
 
     Example:
         >>> gp = SaasFullyBayesianSingleTaskGP(train_X, train_Y)
@@ -372,18 +368,20 @@ def fit_fully_bayesian_model_nuts(
     # Do inference with NUTS
     nuts = NUTS(
         model.pyro_model.sample,
-        dense_mass=True,
+        jit_compile=jit_compile,
+        full_mass=True,
+        ignore_jit_warnings=True,
         max_tree_depth=max_tree_depth,
     )
     mcmc = MCMC(
         nuts,
-        num_warmup=warmup_steps,
+        warmup_steps=warmup_steps,
         num_samples=num_samples,
-        progress_bar=not disable_progbar,
+        disable_progbar=disable_progbar,
     )
-    mcmc.run(jax.random.PRNGKey(seed))
+    mcmc.run()
 
-    # Get final MCMC samples from the NumPyro model
+    # Get final MCMC samples from the Pyro model
     mcmc_samples = model.pyro_model.postprocess_mcmc_samples(
         mcmc_samples=mcmc.get_samples()
     )