Move get_seeds_per_chain and Random type-hint variables to util.py

Author: ricardoV94

.github/workflows/tests.yml

@@ -149,7 +149,7 @@ jobs:
         python-version: ["3.8"]
         test-subset:
           - pymc/tests/variational/test_approximations.py pymc/tests/variational/test_callbacks.py pymc/tests/variational/test_inference.py pymc/tests/variational/test_opvi.py pymc/tests/test_initial_point.py
-          - pymc/tests/test_model.py pymc/tests/test_sampling_utils.py pymc/tests/step_methods/test_compound.py pymc/tests/step_methods/hmc/test_hmc.py
+          - pymc/tests/test_model.py pymc/tests/step_methods/test_compound.py pymc/tests/step_methods/hmc/test_hmc.py
           - pymc/tests/gp/test_cov.py pymc/tests/gp/test_gp.py pymc/tests/gp/test_mean.py pymc/tests/gp/test_util.py pymc/tests/ode/test_ode.py pymc/tests/ode/test_utils.py pymc/tests/smc/test_smc.py pymc/tests/test_parallel_sampling.py
           - pymc/tests/test_sampling.py pymc/tests/step_methods/test_metropolis.py pymc/tests/step_methods/test_slicer.py pymc/tests/step_methods/hmc/test_nuts.py
 

pymc/parallel_sampling.py

@@ -21,7 +21,7 @@
 import traceback
 
 from collections import namedtuple
-from typing import TYPE_CHECKING, Dict, List, Sequence
+from typing import Dict, List, Sequence
 
 import cloudpickle
 import numpy as np
@@ -31,10 +31,7 @@
 from pymc import aesaraf
 from pymc.blocking import DictToArrayBijection
 from pymc.exceptions import SamplingError
-
-# Avoid circular import
-if TYPE_CHECKING:
-    from pymc.sampling import RandomSeed
+from pymc.util import RandomSeed
 
 logger = logging.getLogger("pymc")
 

pymc/sampling.py

@@ -46,12 +46,18 @@
 )
 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
 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
 from pymc.step_methods.hmc import quadpotential
-from pymc.util import drop_warning_stat, get_untransformed_name, is_transformed_name
+from pymc.util import (
+    RandomSeed,
+    RandomState,
+    _get_seeds_per_chain,
+    drop_warning_stat,
+    get_untransformed_name,
+    is_transformed_name,
+)
 from pymc.vartypes import discrete_types
 
 sys.setrecursionlimit(10000)

pymc/sampling_forward.py

@@ -53,8 +53,13 @@
 from pymc.backends.base import MultiTrace
 from pymc.blocking import PointType
 from pymc.model import Model, modelcontext
-from pymc.sampling_utils import RandomState, _get_seeds_per_chain
-from pymc.util import dataset_to_point_list, get_default_varnames, point_wrapper
+from pymc.util import (
+    RandomState,
+    _get_seeds_per_chain,
+    dataset_to_point_list,
+    get_default_varnames,
+    point_wrapper,
+)
 
 __all__ = (
     "compile_forward_sampling_function",

pymc/sampling_jax.py

@@ -8,7 +8,6 @@
 
 from pymc.initial_point import StartDict
 from pymc.sampling import _init_jitter
-from pymc.sampling_utils import RandomSeed, _get_seeds_per_chain
 
 xla_flags = os.getenv("XLA_FLAGS", "")
 xla_flags = re.sub(r"--xla_force_host_platform_device_count=.+\s", "", xla_flags).split()
@@ -33,7 +32,7 @@
 
 from pymc import Model, modelcontext
 from pymc.backends.arviz import find_constants, find_observations
-from pymc.util import get_default_varnames
+from pymc.util import RandomSeed, _get_seeds_per_chain, get_default_varnames
 
 warnings.warn("This module is experimental.")
 

pymc/sampling_utils.py

@@ -14,9 +14,7 @@
 
 """Helper functions for MCMC, prior and posterior predictive sampling."""
 
-from typing import Optional, Sequence, Union
-
-import numpy as np
+from typing import Union
 
 from typing_extensions import TypeAlias
 
@@ -25,58 +23,4 @@
 
 Backend: TypeAlias = Union[BaseTrace, MultiTrace, NDArray]
 
-RandomSeed = Optional[Union[int, Sequence[int], np.ndarray]]
-RandomState = Union[RandomSeed, np.random.RandomState, np.random.Generator]
-
-
 __all__ = ()
-
-
-def _get_seeds_per_chain(
-    random_state: RandomState,
-    chains: int,
-) -> Union[Sequence[int], np.ndarray]:
-    """Obtain or validate specified integer seeds per chain.
-
-    This function process different possible sources of seeding and returns one integer
-    seed per chain:
-    1. If the input is an integer and a single chain is requested, the input is
-        returned inside a tuple.
-    2. If the input is a sequence or NumPy array with as many entries as chains,
-        the input is returned.
-    3. If the input is an integer and multiple chains are requested, new unique seeds
-        are generated from NumPy default Generator seeded with that integer.
-    4. If the input is None new unique seeds are generated from an unseeded NumPy default
-        Generator.
-    5. If a RandomState or Generator is provided, new unique seeds are generated from it.
-
-    Raises
-    ------
-    ValueError
-        If none of the conditions above are met
-    """
-
-    def _get_unique_seeds_per_chain(integers_fn):
-        seeds = []
-        while len(set(seeds)) != chains:
-            seeds = [int(seed) for seed in integers_fn(2**30, dtype=np.int64, size=chains)]
-        return seeds
-
-    if random_state is None or isinstance(random_state, int):
-        if chains == 1 and isinstance(random_state, int):
-            return (random_state,)
-        return _get_unique_seeds_per_chain(np.random.default_rng(random_state).integers)
-    if isinstance(random_state, np.random.Generator):
-        return _get_unique_seeds_per_chain(random_state.integers)
-    if isinstance(random_state, np.random.RandomState):
-        return _get_unique_seeds_per_chain(random_state.randint)
-
-    if not isinstance(random_state, (list, tuple, np.ndarray)):
-        raise ValueError(f"The `seeds` must be array-like. Got {type(random_state)} instead.")
-
-    if len(random_state) != chains:
-        raise ValueError(
-            f"Number of seeds ({len(random_state)}) does not match the number of chains ({chains})."
-        )
-
-    return random_state

pymc/tests/test_sampling_utils.py

@@ -11,6 +11,7 @@
 #   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.
+import re
 
 import arviz
 import numpy as np
@@ -24,6 +25,7 @@
 from pymc.distributions.transforms import RVTransform
 from pymc.util import (
     UNSET,
+    _get_seeds_per_chain,
     dataset_to_point_list,
     drop_warning_stat,
     hash_key,
@@ -198,3 +200,39 @@ def test_drop_warning_stat():
         assert "a" in ss
         assert "warning" not in ss
         assert "warning_dim_0" not in ss
+
+
+def test_get_seeds_per_chain():
+    ret = _get_seeds_per_chain(None, chains=1)
+    assert len(ret) == 1 and isinstance(ret[0], int)
+
+    ret = _get_seeds_per_chain(None, chains=2)
+    assert len(ret) == 2 and isinstance(ret[0], int)
+
+    ret = _get_seeds_per_chain(5, chains=1)
+    assert ret == (5,)
+
+    ret = _get_seeds_per_chain(5, chains=3)
+    assert len(ret) == 3 and isinstance(ret[0], int) and not any(r == 5 for r in ret)
+
+    rng = np.random.default_rng(123)
+    expected_ret = rng.integers(2**30, dtype=np.int64, size=1)
+    rng = np.random.default_rng(123)
+    ret = _get_seeds_per_chain(rng, chains=1)
+    assert ret == expected_ret
+
+    rng = np.random.RandomState(456)
+    expected_ret = rng.randint(2**30, dtype=np.int64, size=2)
+    rng = np.random.RandomState(456)
+    ret = _get_seeds_per_chain(rng, chains=2)
+    assert np.all(ret == expected_ret)
+
+    for expected_ret in ([0, 1, 2], (0, 1, 2, 3), np.arange(5)):
+        ret = _get_seeds_per_chain(expected_ret, chains=len(expected_ret))
+        assert ret is expected_ret
+
+        with pytest.raises(ValueError, match="does not match the number of chains"):
+            _get_seeds_per_chain(expected_ret, chains=len(expected_ret) + 1)
+
+    with pytest.raises(ValueError, match=re.escape("The `seeds` must be array-like")):
+        _get_seeds_per_chain({1: 1, 2: 2}, 2)

pymc/tests/test_util.py

@@ -14,7 +14,7 @@
 
 import functools
 
-from typing import Any, Dict, List, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
 
 import arviz
 import cloudpickle
@@ -387,3 +387,57 @@ def wrapped(**kwargs):
         return core_function(**input_point)
 
     return wrapped
+
+
+RandomSeed = Optional[Union[int, Sequence[int], np.ndarray]]
+RandomState = Union[RandomSeed, np.random.RandomState, np.random.Generator]
+
+
+def _get_seeds_per_chain(
+    random_state: RandomState,
+    chains: int,
+) -> Union[Sequence[int], np.ndarray]:
+    """Obtain or validate specified integer seeds per chain.
+
+    This function process different possible sources of seeding and returns one integer
+    seed per chain:
+    1. If the input is an integer and a single chain is requested, the input is
+        returned inside a tuple.
+    2. If the input is a sequence or NumPy array with as many entries as chains,
+        the input is returned.
+    3. If the input is an integer and multiple chains are requested, new unique seeds
+        are generated from NumPy default Generator seeded with that integer.
+    4. If the input is None new unique seeds are generated from an unseeded NumPy default
+        Generator.
+    5. If a RandomState or Generator is provided, new unique seeds are generated from it.
+
+    Raises
+    ------
+    ValueError
+        If none of the conditions above are met
+    """
+
+    def _get_unique_seeds_per_chain(integers_fn):
+        seeds = []
+        while len(set(seeds)) != chains:
+            seeds = [int(seed) for seed in integers_fn(2**30, dtype=np.int64, size=chains)]
+        return seeds
+
+    if random_state is None or isinstance(random_state, int):
+        if chains == 1 and isinstance(random_state, int):
+            return (random_state,)
+        return _get_unique_seeds_per_chain(np.random.default_rng(random_state).integers)
+    if isinstance(random_state, np.random.Generator):
+        return _get_unique_seeds_per_chain(random_state.integers)
+    if isinstance(random_state, np.random.RandomState):
+        return _get_unique_seeds_per_chain(random_state.randint)
+
+    if not isinstance(random_state, (list, tuple, np.ndarray)):
+        raise ValueError(f"The `seeds` must be array-like. Got {type(random_state)} instead.")
+
+    if len(random_state) != chains:
+        raise ValueError(
+            f"Number of seeds ({len(random_state)}) does not match the number of chains ({chains})."
+        )
+
+    return random_state

pymc/util.py

@@ -70,8 +70,12 @@
 from pymc.blocking import DictToArrayBijection
 from pymc.initial_point import make_initial_point_fn
 from pymc.model import modelcontext
-from pymc.sampling_utils import RandomState, _get_seeds_per_chain
-from pymc.util import WithMemoization, locally_cachedmethod
+from pymc.util import (
+    RandomState,
+    WithMemoization,
+    _get_seeds_per_chain,
+    locally_cachedmethod,
+)
 from pymc.variational.updates import adagrad_window
 from pymc.vartypes import discrete_types