use standard processbar kwarg

Author: junpenglao

pymc/sampling/jax.py

@@ -342,6 +342,7 @@ def sample_blackjax_nuts(
     initvals: Optional[Union[StartDict, Sequence[Optional[StartDict]]]] = None,
     model: Optional[Model] = None,
     var_names: Optional[Sequence[str]] = None,
+    progress_bar: bool = False,
     keep_untransformed: bool = False,
     chain_method: str = "parallel",
     postprocessing_backend: Optional[Literal["cpu", "gpu"]] = None,
@@ -447,21 +448,22 @@ def sample_blackjax_nuts(
     # Adapted from numpyro
     if chain_method == "parallel":
         map_fn = jax.pmap
-        if adaptation_kwargs.get("progress_bar", False):
+        if progress_bar:
             import warnings
 
             warnings.warn(
-                "BlackJax currently only display progress_bar correctly under "
-                "`chain_method == 'vectorized'`. Setting `progress_bar=False`."
+                "BlackJax currently only display progress bar correctly under "
+                "`chain_method == 'vectorized'`. Setting `progressbar=False`."
             )
-            adaptation_kwargs["progress_bar"] = False
+            progress_bar = False
     elif chain_method == "vectorized":
         map_fn = jax.vmap
     else:
         raise ValueError(
             "Only supporting the following methods to draw chains:" ' "parallel" or "vectorized"'
         )
 
+    adaptation_kwargs["progress_bar"] = progress_bar
     get_posterior_samples = partial(
         _blackjax_inference_loop,
         logprob_fn=logprob_fn,

pymc/sampling/mcmc.py

@@ -372,6 +372,7 @@ def _sample_external_nuts(
             random_seed=random_seed,
             initvals=initvals,
             model=model,
+            progress_bar=progressbar,
             idata_kwargs=idata_kwargs,
             **nuts_sampler_kwargs,
         )