Second iteration

Author: ricardoV94

pymc/sampling/external/__init__.py

@@ -11,6 +11,5 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-from pymc.sampling.external.base import ExternalSampler
 from pymc.sampling.external.jax import Blackjax, Numpyro
 from pymc.sampling.external.nutpie import Nutpie

pymc/sampling/external/base.py

@@ -12,37 +12,40 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import Any
+
+from pytensor.scalar import discrete_dtypes
 
 from pymc.model.core import modelcontext
-from pymc.util import get_value_vars_from_user_vars
+from pymc.util import RandomSeed
 
 
 class ExternalSampler(ABC):
-    def __init__(self, vars=None, model=None):
+    def __init__(self, model=None):
         model = modelcontext(model)
-        if vars is None:
-            vars = model.free_RVs
-        else:
-            vars = get_value_vars_from_user_vars(vars, model=model)
-            if set(vars) != set(model.free_RVs):
-                raise ValueError(
-                    "External samplers must sample all the model free_RVs, not just a subset"
-                )
-        self.vars = vars
         self.model = model
 
     @abstractmethod
     def sample(
         self,
-        tune,
-        draws,
-        chains,
-        initvals,
-        random_seed,
-        progressbar,
-        var_names,
-        idata_kwargs,
-        compute_convergence_checks,
+        *,
+        tune: int,
+        draws: int,
+        chains: int,
+        initvals: dict[str, Any] | Sequence[dict[str, Any]],
+        random_seed: RandomSeed,
+        progressbar: bool,
+        var_names: Sequence[str] | None = None,
+        idata_kwargs: dict[str, Any] | None = None,
+        compute_convergence_checks: bool,
         **kwargs,
     ):
         pass
+
+
+class NUTSExternalSampler(ExternalSampler):
+    def __init__(self, model=None):
+        super().__init__(model)
+        if any(var.dtype in discrete_dtypes for var in model.free_RVs):
+            raise ValueError("External NUTS samplers can only sample continuous variables")

pymc/sampling/external/jax.py

@@ -16,24 +16,23 @@
 
 from arviz import InferenceData
 
-from pymc.sampling.external.base import ExternalSampler
+from pymc.sampling.external.base import NUTSExternalSampler
 from pymc.util import RandomState
 
 
-class JAXSampler(ExternalSampler):
-    nuts_sampler = None  # Should be defined by subclass
+class JAXNUTSSampler(NUTSExternalSampler):
+    nuts_sampler: Literal["numpyro", "blackjax"]
 
     def __init__(
         self,
-        vars=None,
         model=None,
         postprocessing_backend: Literal["cpu", "gpu"] | None = None,
         chain_method: Literal["parallel", "vectorized"] = "parallel",
         jitter: bool = True,
         keep_untransformed: bool = False,
         nuts_kwargs: dict | None = None,
     ):
-        super().__init__(vars, model)
+        super().__init__(model)
         self.postprocessing_backend = postprocessing_backend
         self.chain_method = chain_method
         self.jitter = jitter
@@ -53,7 +52,6 @@ def sample(
         idata_kwargs: dict | None = None,
         compute_convergence_checks: bool = True,
         target_accept: float = 0.8,
-        nuts_sampler,
         **kwargs,
     ) -> InferenceData:
         from pymc.sampling.jax import sample_jax_nuts
@@ -80,9 +78,9 @@ def sample(
         )
 
 
-class Numpyro(JAXSampler):
+class Numpyro(JAXNUTSSampler):
     nuts_sampler = "numpyro"
 
 
-class Blackjax(JAXSampler):
+class Blackjax(JAXNUTSSampler):
     nuts_sampler = "blackjax"

pymc/sampling/external/nutpie.py

@@ -14,31 +14,28 @@
 import warnings
 
 from arviz import InferenceData, dict_to_dataset
-from pytensor.scalar import discrete_dtypes
 
 from pymc.backends.arviz import coords_and_dims_for_inferencedata, find_constants, find_observations
-from pymc.sampling.external.base import ExternalSampler
+from pymc.sampling.external.base import NUTSExternalSampler
 from pymc.stats.convergence import log_warnings, run_convergence_checks
 from pymc.util import _get_seeds_per_chain
 
 
-class Nutpie(ExternalSampler):
+class Nutpie(NUTSExternalSampler):
     def __init__(
         self,
-        vars=None,
         model=None,
         backend="numba",
         gradient_backend="pytensor",
         compile_kwargs=None,
         sample_kwargs=None,
     ):
-        super().__init__(vars, model)
-        if any(var.dtype in discrete_dtypes for var in self.vars):
-            raise ValueError("Nutpie can only sample continuous variables")
+        super().__init__(model)
         self.backend = backend
         self.gradient_backend = gradient_backend
         self.compile_kwargs = compile_kwargs or {}
         self.sample_kwargs = sample_kwargs or {}
+        self.compiled_model = None
 
     def sample(
         self,

pymc/sampling/jax.py

@@ -665,7 +665,6 @@ def sample_jax_nuts(
     if "dims" in idata_kwargs:
         dims.update(idata_kwargs.pop("dims"))
 
-    # Use 'partial' to set default arguments before passing 'idata_kwargs'
     idata = az.from_dict(
         posterior=mcmc_samples,
         log_likelihood=log_likelihood,

pymc/sampling/mcmc.py

@@ -292,11 +292,11 @@ def sample(
     chains: int | None = None,
     cores: int | None = None,
     random_seed: RandomState = None,
+    step=None,
+    external_sampler: ExternalSampler | None = None,
     progressbar: bool | ProgressBarType = True,
     progressbar_theme: Theme | None = default_progress_theme,
-    step=None,
     var_names: Sequence[str] | None = None,
-    nuts_sampler: Literal["pymc", "nutpie", "numpyro", "blackjax"] = "pymc",
     initvals: StartDict | Sequence[StartDict | None] | None = None,
     init: str = "auto",
     jitter_max_retries: int = 10,
@@ -324,11 +324,11 @@ def sample(
     chains: int | None = None,
     cores: int | None = None,
     random_seed: RandomState = None,
+    step=None,
+    external_sampler: ExternalSampler | None = None,
     progressbar: bool | ProgressBarType = True,
     progressbar_theme: Theme | None = default_progress_theme,
-    step=None,
     var_names: Sequence[str] | None = None,
-    nuts_sampler: Literal["pymc", "nutpie", "numpyro", "blackjax"] = "pymc",
     initvals: StartDict | Sequence[StartDict | None] | None = None,
     init: str = "auto",
     jitter_max_retries: int = 10,
@@ -356,11 +356,11 @@ def sample(
     chains: int | None = None,
     cores: int | None = None,
     random_seed: RandomState = None,
+    step=None,
+    external_sampler: ExternalSampler | None = None,
     progressbar: bool | ProgressBarType = True,
     progressbar_theme: Theme | None = None,
-    step=None,
     var_names: Sequence[str] | None = None,
-    nuts_sampler: None | Literal["pymc", "nutpie", "numpyro", "blackjax"] = None,
     initvals: StartDict | Sequence[StartDict | None] | None = None,
     init: str = "auto",
     jitter_max_retries: int = 10,
@@ -407,6 +407,12 @@ def sample(
         A ``TypeError`` will be raised if a legacy :py:class:`~numpy.random.RandomState` object is passed.
         We no longer support ``RandomState`` objects because their seeding mechanism does not allow
         easy spawning of new independent random streams that are needed by the step methods.
+    step : function or iterable of functions, optional
+        A step function or collection of functions. If there are variables without step methods,
+        step methods for those variables will be assigned automatically. By default the NUTS step
+        method will be used, if appropriate to the model. Not compatible with external_sampler
+    external_sampler: ExternalSampler, optional
+        An external sampler to sample the whole model. Not compatible with step.
     progressbar: bool or ProgressType, optional
             How and whether to display the progress bar. If False, no progress bar is displayed. Otherwise, you can ask
             for one of the following:
@@ -419,10 +425,6 @@ def sample(
                 are also displayed.
 
             If True, the default is "split+stats" is used.
-    step : function or iterable of functions
-        A step function or collection of functions. If there are variables without step methods,
-        step methods for those variables will be assigned automatically. By default the NUTS step
-        method will be used, if appropriate to the model.
     var_names : list of str, optional
         Names of variables to be stored in the trace. Defaults to all free variables and deterministics.
     nuts_sampler : str
@@ -608,35 +610,38 @@ def joined_blas_limiter():
     rngs = get_random_generator(random_seed).spawn(chains)
     random_seed_list = [rng.integers(2**30) for rng in rngs]
 
-    if step is None and nuts_sampler not in (None, "pymc"):
-        # Temporarily instantiate external samplers for user, for backwards-compat
+    if "nuts_sampler" in kwargs:
+        # Transition backwards-compatibility
+        nuts_sampler = kwargs.pop("nuts_sampler")
         warnings.warn(
             f"Setting `pm.sample(nuts_sampler='{nuts_sampler}, nuts_sampler_kwargs=...)'` is deprecated.\n"
-            f"Use `pm.sample(step=pm.external.{nuts_sampler.capitalize()}(**nuts_sampler_kwargs))` instead",
+            f"Use `pm.sample(external_sampler=pm.external.{nuts_sampler.capitalize()}(**nuts_sampler_kwargs))` instead",
             FutureWarning,
         )
         from pymc.sampling import external
 
-        step = getattr(external, nuts_sampler.capitalize())(
+        external_sampler = getattr(external, nuts_sampler.capitalize())(
             model=model,
             **(nuts_sampler_kwargs or {}),
         )
         nuts_sampler_kwargs = None
 
-    if isinstance(step, list | tuple) and len(step) == 1:
-        [step] = step
+    if external_sampler is not None:
+        if step is not None:
+            raise ValueError("`step` and `external_sampler` cannot be used together")
 
-    if isinstance(step, ExternalSampler):
-        if step.model is not model:
-            raise ValueError("External step model does not match model detected by sample")
+        if external_sampler.model is not model:
+            raise ValueError(
+                "External sampler model does not match model detected by sample function"
+            )
         if nuts_sampler_kwargs:
             raise ValueError(
                 f"{nuts_sampler_kwargs=} should be passed when constructing external sampler"
             )
         if "nuts" in kwargs:
-            kwargs.update(kwargs["nuts"].pop())
+            kwargs.update(kwargs.pop("nuts"))
         with joined_blas_limiter():
-            return step.sample(
+            return external_sampler.sample(
                 tune=tune,
                 draws=draws,
                 chains=chains,

tests/sampling/test_jax.py

@@ -82,8 +82,7 @@ def test_jax_PosDefMatrix():
         ),
     ],
 )
-@pytest.mark.parametrize("postprocessing_vectorize", ["scan", "vmap"])
-def test_transform_samples(sampler, postprocessing_backend, chains, postprocessing_vectorize):
+def test_transform_samples(sampler, postprocessing_backend, chains):
     pytensor.config.on_opt_error = "raise"
     np.random.seed(13244)
 
@@ -99,7 +98,6 @@ def test_transform_samples(sampler, postprocessing_backend, chains, postprocessi
             random_seed=1322,
             keep_untransformed=True,
             postprocessing_backend=postprocessing_backend,
-            postprocessing_vectorize=postprocessing_vectorize,
         )
 
     log_vals = trace.posterior["sigma_log__"].values

tests/sampling/test_mcmc_external.py

@@ -59,8 +59,7 @@ def test_external_nuts_sampler(recwarn, nuts_sampler):
         expected.add(
             (
                 UserWarning,
-                "`initvals` are currently not passed to nutpie sampler. "
-                "Use `init_mean` kwarg following nutpie specification instead.",
+                "initvals are currently ignored by the nutpie sampler.",
             )
         )
     assert warns == expected