Add chain_method kwarg to sample_numpyro_nuts

Default behaviour is unchanged, but setting `chain_method='vectorized'` can be
more efficient when run on a single GPU.

Author: martiningram

pymc/sampling_jax.py

@@ -141,6 +141,7 @@ def sample_numpyro_nuts(
     var_names=None,
     progress_bar=True,
     keep_untransformed=False,
+    chain_method="parallel",
 ):
     from numpyro.infer import MCMC, NUTS
 
@@ -188,7 +189,7 @@ def sample_numpyro_nuts(
         num_samples=draws,
         num_chains=chains,
         postprocess_fn=None,
-        chain_method="parallel",
+        chain_method=chain_method,
         progress_bar=progress_bar,
     )