Refactor RandomState/Generator reseeding logic in initial_point into aesaraf

Author: ricardoV94

pymc/aesaraf.py

@@ -20,6 +20,7 @@
     Iterable,
     List,
     Optional,
+    Sequence,
     Set,
     Tuple,
     Union,
@@ -893,6 +894,40 @@ def local_check_parameter_to_ninf_switch(fgraph, node):
 )
 
 
+def find_rng_nodes(variables: Iterable[TensorVariable]):
+    """Return RNG variables in a graph"""
+    return [
+        node
+        for node in graph_inputs(variables)
+        if isinstance(
+            node,
+            (
+                at.random.var.RandomStateSharedVariable,
+                at.random.var.RandomGeneratorSharedVariable,
+            ),
+        )
+    ]
+
+
+SeedSequenceSeed = Optional[Union[int, Sequence[int], np.ndarray, np.random.SeedSequence]]
+
+
+def reseed_rngs(
+    rngs: Sequence[SharedVariable],
+    seed: SeedSequenceSeed,
+) -> None:
+    """Create a new set of RandomState/Generator for each rng based on a seed"""
+    bit_generators = [
+        np.random.PCG64(sub_seed) for sub_seed in np.random.SeedSequence(seed).spawn(len(rngs))
+    ]
+    for rng, bit_generator in zip(rngs, bit_generators):
+        if isinstance(rng, at.random.var.RandomStateSharedVariable):
+            new_rng = np.random.RandomState(bit_generator)
+        else:
+            new_rng = np.random.Generator(bit_generator)
+        rng.set_value(new_rng, borrow=True)
+
+
 def compile_pymc(
     inputs, outputs, mode=None, **kwargs
 ) -> Callable[..., Union[np.ndarray, List[np.ndarray]]]:

pymc/initial_point.py

@@ -20,11 +20,11 @@
 import aesara.tensor as at
 import numpy as np
 
-from aesara.graph.basic import Variable, graph_inputs
+from aesara.graph.basic import Variable
 from aesara.graph.fg import FunctionGraph
 from aesara.tensor.var import TensorVariable
 
-from pymc.aesaraf import compile_pymc
+from pymc.aesaraf import compile_pymc, find_rng_nodes, reseed_rngs
 from pymc.util import get_transformed_name, get_untransformed_name, is_transformed_name
 
 StartDict = Dict[Union[Variable, str], Union[np.ndarray, Variable, str]]
@@ -150,19 +150,6 @@ def make_initial_point_fn(
         If `True` the returned variables will correspond to transformed initial values.
     """
 
-    def find_rng_nodes(variables):
-        return [
-            node
-            for node in graph_inputs(variables)
-            if isinstance(
-                node,
-                (
-                    at.random.var.RandomStateSharedVariable,
-                    at.random.var.RandomGeneratorSharedVariable,
-                ),
-            )
-        ]
-
     sdict_overrides = convert_str_to_rv_dict(model, overrides or {})
     initval_strats = {
         **model.initial_values,
@@ -208,16 +195,7 @@ def make_seeded_function(func):
 
         @functools.wraps(func)
         def inner(seed, *args, **kwargs):
-            seeds = [
-                np.random.PCG64(sub_seed)
-                for sub_seed in np.random.SeedSequence(seed).spawn(len(rngs))
-            ]
-            for rng, seed in zip(rngs, seeds):
-                if isinstance(rng, at.random.var.RandomStateSharedVariable):
-                    new_rng = np.random.RandomState(seed)
-                else:
-                    new_rng = np.random.Generator(seed)
-                rng.set_value(new_rng, True)
+            reseed_rngs(rngs, seed)
             values = func(*args, **kwargs)
             return dict(zip(varnames, values))
 

pymc/tests/test_aesaraf.py

@@ -26,6 +26,7 @@
 from aesara.graph.basic import Constant, Variable, ancestors, equal_computations
 from aesara.tensor.random.basic import normal, uniform
 from aesara.tensor.random.op import RandomVariable
+from aesara.tensor.random.var import RandomStateSharedVariable
 from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
 from aesara.tensor.var import TensorVariable
 
@@ -36,6 +37,7 @@
     compile_pymc,
     convert_observed_data,
     extract_obs_data,
+    reseed_rngs,
     rvs_to_value_vars,
     walk_model,
 )
@@ -507,3 +509,35 @@ def update(self, node):
         fn = compile_pymc(inputs=[], outputs=dummy_x)
         assert fn() == 2.0
         assert fn() == 3.0
+
+
+def test_reseed_rngs():
+    # Reseed_rngs uses the `PCG64` bit_generator, which is currently the default
+    # bit_generator used by NumPy. If this default changes in the future, this test will
+    # catch that. We will then have to decide whether to switch to the new default in
+    # PyMC or whether to stick with the older one (PCG64). This will pose a trade-off
+    # between backwards reproducibility and better/faster seeding. If we decide to change,
+    # the next line should be updated:
+    default_rng = np.random.PCG64
+    assert isinstance(np.random.default_rng().bit_generator, default_rng)
+
+    seed = 543
+
+    bit_generators = [default_rng(sub_seed) for sub_seed in np.random.SeedSequence(seed).spawn(2)]
+
+    rngs = [
+        aesara.shared(rng_type(default_rng()))
+        for rng_type in (np.random.Generator, np.random.RandomState)
+    ]
+    for rng, bit_generator in zip(rngs, bit_generators):
+        if isinstance(rng, RandomStateSharedVariable):
+            assert rng.get_value()._bit_generator.state != bit_generator.state
+        else:
+            assert rng.get_value().bit_generator.state != bit_generator.state
+
+    reseed_rngs(rngs, seed)
+    for rng, bit_generator in zip(rngs, bit_generators):
+        if isinstance(rng, RandomStateSharedVariable):
+            assert rng.get_value()._bit_generator.state == bit_generator.state
+        else:
+            assert rng.get_value().bit_generator.state == bit_generator.state