Move non-shared utils to respective files

Author: ricardoV94

pymc/sampling.py

@@ -46,12 +46,7 @@
 )
 from pymc.model import Model, modelcontext
 from pymc.parallel_sampling import Draw, _cpu_count
-from pymc.sampling_utils import (
-    RandomSeed,
-    RandomState,
-    _get_seeds_per_chain,
-    all_continuous,
-)
+from pymc.sampling_utils import RandomSeed, RandomState, _get_seeds_per_chain
 from pymc.stats.convergence import SamplerWarning, log_warning, run_convergence_checks
 from pymc.step_methods import NUTS, CompoundStep, DEMetropolis
 from pymc.step_methods.arraystep import BlockedStep, PopulationArrayStepShared
@@ -208,6 +203,17 @@ def _print_step_hierarchy(s: Step, level: int = 0) -> None:
         _log.info(">" * level + f"{s.__class__.__name__}: [{varnames}]")
 
 
+def all_continuous(vars):
+    """Check that vars not include discrete variables, excepting observed RVs."""
+
+    vars_ = [var for var in vars if not hasattr(var.tag, "observations")]
+
+    if any([(var.dtype in discrete_types) for var in vars_]):
+        return False
+    else:
+        return True
+
+
 def sample(
     draws: int = 1000,
     step=None,

pymc/sampling_forward.py

@@ -44,20 +44,16 @@
 from aesara.tensor.sharedvar import SharedVariable
 from arviz import InferenceData
 from fastprogress.fastprogress import progress_bar
+from typing_extensions import TypeAlias
 
 import pymc as pm
 
 from pymc.aesaraf import compile_pymc
 from pymc.backends.arviz import _DefaultTrace
 from pymc.backends.base import MultiTrace
+from pymc.blocking import PointType
 from pymc.model import Model, modelcontext
-from pymc.sampling_utils import (
-    ArrayLike,
-    PointList,
-    RandomState,
-    _get_seeds_per_chain,
-    get_vars_in_point_list,
-)
+from pymc.sampling_utils import RandomState, _get_seeds_per_chain
 from pymc.util import dataset_to_point_list, get_default_varnames, point_wrapper
 
 __all__ = (
@@ -69,9 +65,22 @@
 )
 
 
+ArrayLike: TypeAlias = Union[np.ndarray, List[float]]
+PointList: TypeAlias = List[PointType]
+
 _log = logging.getLogger("pymc")
 
 
+def get_vars_in_point_list(trace, model):
+    """Get the list of Variable instances in the model that have values stored in the trace."""
+    if not isinstance(trace, MultiTrace):
+        names_in_trace = list(trace[0])
+    else:
+        names_in_trace = trace.varnames
+    vars_in_trace = [model[v] for v in names_in_trace if v in model]
+    return vars_in_trace
+
+
 def compile_forward_sampling_function(
     outputs: List[Variable],
     vars_in_trace: List[Variable],

pymc/sampling_utils.py

@@ -14,19 +14,15 @@
 
 """Helper functions for MCMC, prior and posterior predictive sampling."""
 
-from typing import List, Optional, Sequence, Union
+from typing import Optional, Sequence, Union
 
 import numpy as np
 
 from typing_extensions import TypeAlias
 
 from pymc.backends.base import BaseTrace, MultiTrace
 from pymc.backends.ndarray import NDArray
-from pymc.initial_point import PointType
-from pymc.vartypes import discrete_types
 
-ArrayLike: TypeAlias = Union[np.ndarray, List[float]]
-PointList: TypeAlias = List[PointType]
 Backend: TypeAlias = Union[BaseTrace, MultiTrace, NDArray]
 
 RandomSeed = Optional[Union[int, Sequence[int], np.ndarray]]
@@ -36,17 +32,6 @@
 __all__ = ()
 
 
-def all_continuous(vars):
-    """Check that vars not include discrete variables, excepting observed RVs."""
-
-    vars_ = [var for var in vars if not hasattr(var.tag, "observations")]
-
-    if any([(var.dtype in discrete_types) for var in vars_]):
-        return False
-    else:
-        return True
-
-
 def _get_seeds_per_chain(
     random_state: RandomState,
     chains: int,
@@ -95,13 +80,3 @@ def _get_unique_seeds_per_chain(integers_fn):
         )
 
     return random_state
-
-
-def get_vars_in_point_list(trace, model):
-    """Get the list of Variable instances in the model that have values stored in the trace."""
-    if not isinstance(trace, MultiTrace):
-        names_in_trace = list(trace[0])
-    else:
-        names_in_trace = trace.varnames
-    vars_in_trace = [model[v] for v in names_in_trace if v in model]
-    return vars_in_trace

pymc/tests/test_sampling_forward.py

@@ -35,7 +35,10 @@
 
 from pymc.aesaraf import compile_pymc
 from pymc.backends.base import MultiTrace
-from pymc.sampling_forward import compile_forward_sampling_function
+from pymc.sampling_forward import (
+    compile_forward_sampling_function,
+    get_vars_in_point_list,
+)
 from pymc.tests.helpers import SeededTest, fast_unstable_sampling_mode
 
 
@@ -1635,3 +1638,24 @@ def test_Triangular(
             prior_samples=prior_samples,
         )
         assert prior["target"].shape == (prior_samples,) + shape
+
+
+def test_get_vars_in_point_list():
+    with pm.Model() as modelA:
+        pm.Normal("a", 0, 1)
+        pm.Normal("b", 0, 1)
+    with pm.Model() as modelB:
+        a = pm.Normal("a", 0, 1)
+        pm.Normal("c", 0, 1)
+
+    point_list = [{"a": 0, "b": 0}]
+    vars_in_trace = get_vars_in_point_list(point_list, modelB)
+    assert set(vars_in_trace) == {a}
+
+    strace = pm.backends.NDArray(model=modelB, vars=modelA.free_RVs)
+    strace.setup(1, 1)
+    strace.values = point_list[0]
+    strace.draw_idx = 1
+    trace = MultiTrace([strace])
+    vars_in_trace = get_vars_in_point_list(trace, modelB)
+    assert set(vars_in_trace) == {a}

pymc/tests/test_sampling_utils.py

@@ -16,10 +16,7 @@
 import numpy as np
 import pytest
 
-import pymc as pm
-
-from pymc.backends.base import MultiTrace
-from pymc.sampling_utils import _get_seeds_per_chain, get_vars_in_point_list
+from pymc.sampling_utils import _get_seeds_per_chain
 
 
 def test_get_seeds_per_chain():
@@ -56,24 +53,3 @@ def test_get_seeds_per_chain():
 
     with pytest.raises(ValueError, match=re.escape("The `seeds` must be array-like")):
         _get_seeds_per_chain({1: 1, 2: 2}, 2)
-
-
-def test_get_vars_in_point_list():
-    with pm.Model() as modelA:
-        pm.Normal("a", 0, 1)
-        pm.Normal("b", 0, 1)
-    with pm.Model() as modelB:
-        a = pm.Normal("a", 0, 1)
-        pm.Normal("c", 0, 1)
-
-    point_list = [{"a": 0, "b": 0}]
-    vars_in_trace = get_vars_in_point_list(point_list, modelB)
-    assert set(vars_in_trace) == {a}
-
-    strace = pm.backends.NDArray(model=modelB, vars=modelA.free_RVs)
-    strace.setup(1, 1)
-    strace.values = point_list[0]
-    strace.draw_idx = 1
-    trace = MultiTrace([strace])
-    vars_in_trace = get_vars_in_point_list(trace, modelB)
-    assert set(vars_in_trace) == {a}