Add support for numpyro and blackjax PyMC samplers

bambi/__init__.py

@@ -1,7 +1,5 @@
 import logging
 
-from pymc import math
-
 from .data import clear_data_home, load_data
 from .families import Family, Likelihood, Link
 from .models import Model

bambi/backend/pymc.py

@@ -87,7 +87,7 @@ def run(
     ):
         """Run PyMC sampler."""
         # NOTE: Methods return different types of objects (idata, approximation, and dictionary)
-        if method.lower() == "mcmc":
+        if method.lower() in ["mcmc", "nuts_numpyro", "nuts_blackjax"]:
             result = self._run_mcmc(
                 draws,
                 tune,
@@ -99,6 +99,7 @@ def run(
                 chains,
                 cores,
                 random_seed,
+                method.lower(),
                 **kwargs,
             )
         elif method.lower() == "vi":
@@ -209,40 +210,69 @@ def _run_mcmc(
         chains=None,
         cores=None,
         random_seed=None,
+        sampler_backend="mcmc",
         **kwargs,
     ):
         with self.model:
-            try:
-                idata = pm.sample(
-                    draws=draws,
-                    tune=tune,
-                    discard_tuned_samples=discard_tuned_samples,
-                    init=init,
-                    n_init=n_init,
-                    chains=chains,
-                    cores=cores,
-                    random_seed=random_seed,
-                    **kwargs,
-                )
-            except (RuntimeError, ValueError):
-                if "ValueError: Mass matrix contains" in traceback.format_exc() and init == "auto":
-                    _log.info(
-                        "\nThe default initialization using init='auto' has failed, trying to "
-                        "recover by switching to init='adapt_diag'",
-                    )
+            if sampler_backend == "mcmc":
+                try:
                     idata = pm.sample(
                         draws=draws,
                         tune=tune,
                         discard_tuned_samples=discard_tuned_samples,
-                        init="adapt_diag",
+                        init=init,
                         n_init=n_init,
                         chains=chains,
                         cores=cores,
                         random_seed=random_seed,
                         **kwargs,
                     )
-                else:
-                    raise
+                except (RuntimeError, ValueError):
+                    if "ValueError: Mass matrix contains" in traceback.format_exc() and init == "auto":
+                        _log.info(
+                            "\nThe default initialization using init='auto' has failed, trying to "
+                            "recover by switching to init='adapt_diag'",
+                        )
+                        idata = pm.sample(
+                            draws=draws,
+                            tune=tune,
+                            discard_tuned_samples=discard_tuned_samples,
+                            init="adapt_diag",
+                            n_init=n_init,
+                            chains=chains,
+                            cores=cores,
+                            random_seed=random_seed,
+                            **kwargs,
+                        )
+                    else:
+                        raise
+            elif sampler_backend == "nuts_numpyro":
+                import pymc.sampling_jax  # Lazy import to not force users to install Jax
+                if not chains:
+                    chains = 4  # sample_numpyro_nuts does not handle chains = None like pm.sample does
+                idata = pm.sampling_jax.sample_numpyro_nuts(
+                    draws=draws,
+                    tune=tune,
+                    chains=chains,
+                    random_seed=random_seed,
+                    **kwargs,
+                )
+            elif sampler_backend == "nuts_blackjax":
+                import pymc.sampling_jax  # Lazy import to not force users to install Jax
+                if not chains:
+                    chains = 4  # sample_blackjax_nuts does not handle chains = None like pm.sample does
+                idata = pm.sampling_jax.sample_blackjax_nuts(
+                    draws=draws,
+                    tune=tune,
+                    chains=chains,
+                    random_seed=random_seed,
+                    **kwargs,
+                )
+            else:
+                raise ValueError(
+                    f'sampler_backend value {sampler_backend} is not valid. Please choose one of'
+                    f'``mcmc``, ``nuts_numpyro`` or ``nuts_blackjax``'
+                )
 
         idata = self._clean_mcmc_results(idata, omit_offsets, include_mean)
         return idata

bambi/models.py

@@ -208,6 +208,11 @@ def fit(
             using the ``fit`` function.
             Finally, ``"laplace"``, in which case a Laplace approximation is used and is not
             recommended other than for pedagogical use.
+            To use the PyMC numpyro and blackjax samplers, use ``nuts_numpyro`` or ``nuts_blackjax``
+            respectively. Both methods will only work if you can use NUTS sampling, so your model must
+            be differentiable. It is also recommended, if using ``nuts_numpyro`` or
+            ``nuts_blackjax``, to set the kwarg of ``chain_method`` to either ``parallel`` or
+            ``vectorized`` for optimal performance.
         init: str
             Initialization method. Defaults to ``"auto"``. The available methods are:
             * auto: Use ``"jitter+adapt_diag"`` and if this method fails it uses ``"adapt_diag"``.