Add support for numpyro and blackjax PyMC samplers

bambi/backend/pymc.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import pymc as pm
+import pymc.sampling_jax
 
 import aesara.tensor as at
 
@@ -83,6 +84,7 @@ def run(
         chains=None,
         cores=None,
         random_seed=None,
+        sampler_backend="default",
         **kwargs,
     ):
         """Run PyMC sampler."""
@@ -99,6 +101,7 @@ def run(
                 chains,
                 cores,
                 random_seed,
+                sampler_backend,
                 **kwargs,
             )
         elif method.lower() == "vi":
@@ -209,40 +212,63 @@ def _run_mcmc(
         chains=None,
         cores=None,
         random_seed=None,
+        sampler_backend="default",
         **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 == "default":
+                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 == "numpyro":
+                idata = pm.sampling_jax.sample_numpyro_nuts(
+                    draws=draws,
+                    tune=tune,
+                    chains=chains,
+                    random_seed=random_seed,
+                    **kwargs,
+                )
+            elif sampler_backend == "blackjax":
+                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'``default``, ``numpyro`` or ``blackjax``'
+                )
 
         idata = self._clean_mcmc_results(idata, omit_offsets, include_mean)
         return idata

bambi/models.py

@@ -180,6 +180,7 @@ def fit(
         chains=None,
         cores=None,
         random_seed=None,
+        sampler_backend="default",
         **kwargs,
     ):
         """Fit the model using PyMC.
@@ -239,6 +240,17 @@ def fit(
             in the system unless there are more than 4 CPUs, in which case it is set to 4.
         random_seed : int or list of ints
             A list is accepted if cores is greater than one.
+        sampler_backend: str
+            If ``default`` uses the standard PyMC fit() method. Is only valid if method == "mcmc".
+            Other valid sampling methods are ``numpyro`` and ``blackjax`` which both use JAX (see
+            https://www.pymc.io/blog/v4_announcement.html#new-jax-backend-for-faster-sampling)
+            and which offers both performance improvements on the CPU due to JIT compilation
+            and also GPU use (which will happen automatically if a GPU is detected - see the
+            JAX documentation https://jax.readthedocs.io/en/latest/index.html). Also, both
+            non-default methods will only work if you can use NUTS sampling, so your model must
+            be differentiable. It is also strongly recommended, if using ``numpyro`` or
+            ``blackjax``, to set the kwarg of ``chain_method`` to either ``parallel`` or
+            ``vectorized`` for optimal performance.
         **kwargs:
             For other kwargs see the documentation for ``PyMC.sample()``.
 
@@ -248,6 +260,11 @@ def fit(
         An ``Approximation`` object if  ``"vi"`` and a dictionary if  ``"laplace"``.
         """
 
+        if sampler_backend != "default" and method != "mcmc":
+            raise ValueError(
+                f"Non-default sampler_backend {sampler_backend} can only be used with method 'mcmc'"
+            )
+
         if not self.built:
             self.build()
 
@@ -271,6 +288,7 @@ def fit(
             chains=chains,
             cores=cores,
             random_seed=random_seed,
+            sampler_backend=sampler_backend,
             **kwargs,
         )