mypy + cleanup

pymc-devs · Jan 25, 2025 · b9b0583 · b9b0583
1 parent 4e535d4
commit b9b0583
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git a/pymc/sampling/mcmc.py b/pymc/sampling/mcmc.py
@@ -302,7 +302,7 @@ def _sample_external_nuts(
     initvals: StartDict | Sequence[StartDict | None] | None,
     model: Model,
     var_names: Sequence[str] | None,
-    progressbar: bool,
+    progressbar: bool | ProgressType,
     idata_kwargs: dict | None,
     compute_convergence_checks: bool,
     nuts_sampler_kwargs: dict | None,
@@ -401,7 +401,7 @@ def _sample_external_nuts(
             initvals=initvals,
             model=model,
             var_names=var_names,
-            progressbar=progressbar,
+            progressbar=True if progressbar else False,
             nuts_sampler=sampler,
             idata_kwargs=idata_kwargs,
             compute_convergence_checks=compute_convergence_checks,
@@ -488,7 +488,7 @@ def sample(
     cores: int | None = None,
     random_seed: RandomState = None,
     progressbar: bool | ProgressType = True,
-    progressbar_theme: Theme | None = default_progress_theme,
+    progressbar_theme: Theme | None = None,
     step=None,
     var_names: Sequence[str] | None = None,
     nuts_sampler: Literal["pymc", "nutpie", "numpyro", "blackjax"] = "pymc",
@@ -831,7 +831,9 @@ def joined_blas_limiter():
                 n_init=n_init,
                 model=model,
                 random_seed=random_seed_list,
-                progressbar=progressbar,
+                progressbar=True
+                if progressbar
+                else False,  # ADVI doesn't use the ProgressManager; pass a bool only
                 jitter_max_retries=jitter_max_retries,
                 tune=tune,
                 initvals=initvals,

diff --git a/pymc/util.py b/pymc/util.py
@@ -47,7 +47,7 @@
 from pymc.exceptions import BlockModelAccessError
 
 if TYPE_CHECKING:
-    from pymc import BlockedStep
+    from pymc.step_methods.compound import BlockedStep, CompoundStep
 
 
 ProgressType = Literal[
@@ -727,12 +727,12 @@ class ProgressManager:
 
     def __init__(
         self,
-        step_method: BlockedStep,
+        step_method: "BlockedStep" | "CompoundStep",
         chains: int,
         draws: int,
         tune: int,
         progressbar: bool | ProgressType = True,
-        progressbar_theme: Theme = default_progress_theme,
+        progressbar_theme: Theme | None = None,
     ):
         """
         Manage progress bars displayed during sampling.
@@ -770,6 +770,9 @@ def __init__(
         progressbar_theme: Theme, optional
             The theme to use for the progress bar. Defaults to the default theme.
         """
+        if progressbar_theme is None:
+            progressbar_theme = default_progress_theme
+
         self.combined_progress = False
         self.full_stats = True
         show_progress = True