fix: pin PyMC version below 5.20.1 to avoid TypeError (closes #1397) (#1697)

* fix: pin PyMC version below 5.20.1 to avoid TypeError (closes #1397)

* fix: Disable PyMC progress bar to resolve TypeError

* Add note and warning  and update tests

Author: Dev-Sudarshan

sbi/inference/posteriors/mcmc_posterior.py

@@ -908,6 +908,16 @@ def _pymc_mcmc(
         num_chains = mp.cpu_count() - 1 if num_chains is None else num_chains
         steps = dict(slice_pymc="slice", hmc_pymc="hmc", nuts_pymc="nuts")
 
+        if show_progress_bars:
+            warn(
+                "Note: progress bars for PyMC sampling are disabled due to an "
+                "incompatibility with PyMC>=5.20.1; "
+                "progressbar will be set to False. See PR #1697: "
+                "https://github.com/sbi-dev/sbi/pull/1697",
+                UserWarning,
+                stacklevel=2,
+            )
+
         sampler = PyMCSampler(
             potential_fn=potential_function,
             step=steps[mcmc_method],
@@ -916,7 +926,7 @@ def _pymc_mcmc(
             tune=warmup_steps,
             chains=num_chains,
             mp_ctx=mp_context,
-            progressbar=show_progress_bars,
+            progressbar=False,
             param_name=self.param_name,
             device=self._device,
         )

sbi/samplers/mcmc/pymc_wrapper.py

@@ -112,7 +112,7 @@ def __init__(
         tune: int = 1000,
         chains: Optional[int] = None,
         mp_ctx: str = "spawn",
-        progressbar: bool = True,
+        progressbar: bool = False,
         param_name: str = "theta",
         device: str = "cpu",
     ):
@@ -127,6 +127,10 @@ def __init__(
             chains: Number of MCMC chains to run in parallel.
             mp_ctx: Multiprocessing context for parallel sampling.
             progressbar: Whether to show/hide progress bars.
+                Note: Progress bars are disabled for PyMC sampling in this
+                project due to an incompatibility from PyMC >= 5.20.1 and
+                that affects progress display.
+                See PR #1697 for details: https://github.com/sbi-dev/sbi/pull/1697
             param_name: Name for parameter variable, for PyMC and ArviZ structures
             device: The device to which to move the parameters for potential_fn.
         """

tests/inference_on_device_test.py

@@ -3,11 +3,9 @@
 
 from __future__ import annotations
 
-import sys
 from dataclasses import asdict
 from typing import Tuple, Union
 
-import pymc
 import pytest
 import torch
 import torch.distributions.transforms as torch_tf
@@ -76,19 +74,7 @@
         pytest.param(NRE_B, "resnet", "slice_np", marks=pytest.mark.mcmc),
         (NRE_C, "resnet", "rejection"),
         (NRE_C, "resnet", "importance"),
-        pytest.param(
-            NRE_C,
-            "resnet",
-            "nuts_pymc",
-            marks=(
-                pytest.mark.mcmc,
-                pytest.mark.skipif(
-                    condition=sys.version_info >= (3, 10)
-                    and pymc.__version__ >= "5.20.1",
-                    reason="Inconsistent behaviour with pymc>=5.20.1 and python>=3.10",
-                ),
-            ),
-        ),
+        pytest.param(NRE_C, "resnet", "nuts_pymc", marks=pytest.mark.mcmc),
     ],
 )
 @pytest.mark.parametrize(

tests/linearGaussian_snle_test.py

@@ -3,10 +3,8 @@
 
 from __future__ import annotations
 
-import sys
 from dataclasses import asdict
 
-import pymc
 import pytest
 import torch
 from torch import eye, ones, zeros
@@ -397,18 +395,7 @@ def simulator(theta):
         pytest.param("slice_np", "uniform", marks=pytest.mark.mcmc),
         pytest.param("slice_np_vectorized", "gaussian", marks=pytest.mark.mcmc),
         pytest.param("slice_np_vectorized", "uniform", marks=pytest.mark.mcmc),
-        pytest.param(
-            "nuts_pymc",
-            "gaussian",
-            marks=(
-                pytest.mark.mcmc,
-                pytest.mark.skipif(
-                    condition=sys.version_info >= (3, 10)
-                    and pymc.__version__ >= "5.20.1",
-                    reason="Inconsistent behaviour with pymc>=5.20.1 and python>=3.10",
-                ),
-            ),
-        ),
+        pytest.param("nuts_pymc", "gaussian", marks=pytest.mark.mcmc),
         pytest.param("nuts_pyro", "uniform", marks=pytest.mark.mcmc),
         pytest.param("hmc_pymc", "gaussian", marks=pytest.mark.mcmc),
         ("rejection", "uniform"),

tests/linearGaussian_snre_test.py

@@ -3,10 +3,8 @@
 
 from __future__ import annotations
 
-import sys
 from dataclasses import asdict
 
-import pymc
 import pytest
 from torch import eye, ones, zeros
 from torch.distributions import MultivariateNormal
@@ -319,18 +317,7 @@ def simulator(theta):
         pytest.param("slice_np", "uniform", marks=pytest.mark.mcmc),
         pytest.param("slice_np_vectorized", "gaussian", marks=pytest.mark.mcmc),
         pytest.param("slice_np_vectorized", "uniform", marks=pytest.mark.mcmc),
-        pytest.param(
-            "nuts_pymc",
-            "gaussian",
-            marks=(
-                pytest.mark.mcmc,
-                pytest.mark.skipif(
-                    condition=sys.version_info >= (3, 10)
-                    and pymc.__version__ >= "5.20.1",
-                    reason="Inconsistent behaviour with pymc>=5.20.1 and python>=3.10",
-                ),
-            ),
-        ),
+        pytest.param("nuts_pymc", "gaussian", marks=pytest.mark.mcmc),
         pytest.param("nuts_pyro", "uniform", marks=pytest.mark.mcmc),
         pytest.param("hmc_pyro", "gaussian", marks=pytest.mark.mcmc),
         ("rejection", "uniform"),

tests/mcmc_test.py

@@ -3,11 +3,9 @@
 
 from __future__ import annotations
 
-import sys
 from dataclasses import asdict
 
 import numpy as np
-import pymc
 import pytest
 import torch
 from torch import eye, ones, zeros
@@ -193,13 +191,7 @@ def lp_f(x, track_gradients=True):
     (
         "nuts_pyro",
         "hmc_pyro",
-        pytest.param(
-            "nuts_pymc",
-            marks=pytest.mark.skipif(
-                condition=sys.version_info >= (3, 10) and pymc.__version__ >= "5.20.1",
-                reason="Inconsistent behaviour with pymc>=5.20.1 and python>=3.10",
-            ),
-        ),
+        "nuts_pymc",
         "hmc_pymc",
         "slice_pymc",
         "slice_np",

tests/posterior_sampler_test.py

@@ -3,10 +3,8 @@
 
 from __future__ import annotations
 
-import sys
 from dataclasses import asdict
 
-import pymc
 import pytest
 from pyro.infer.mcmc import MCMC
 from torch import Tensor, eye, zeros
@@ -30,13 +28,7 @@
         "slice_np_vectorized",
         "nuts_pyro",
         "hmc_pyro",
-        pytest.param(
-            "nuts_pymc",
-            marks=pytest.mark.skipif(
-                condition=sys.version_info >= (3, 10) and pymc.__version__ >= "5.20.1",
-                reason="Inconsistent behaviour with pymc>=5.20.1 and python>=3.10",
-            ),
-        ),
+        "nuts_pymc",
         "hmc_pymc",
         "slice_pymc",
     ),