Incrementally update the RNG state in Model

Author: brandonwillard

pymc3/distributions/distribution.py

@@ -137,7 +137,7 @@ def __new__(cls, name, *args, **kwargs):
         rng = kwargs.pop("rng", None)
 
         if rng is None:
-            rng = model.default_rng
+            rng = model.next_rng()
 
         if not isinstance(name, string_types):
             raise TypeError(f"Name needs to be a string but got: {name}")

pymc3/model.py

@@ -42,6 +42,7 @@
 from aesara.graph.basic import Constant, Variable, graph_inputs
 from aesara.graph.fg import FunctionGraph, MissingInputError
 from aesara.tensor.random.opt import local_subtensor_rv_lift
+from aesara.tensor.random.var import RandomStateSharedVariable
 from aesara.tensor.sharedvar import ScalarSharedVariable
 from aesara.tensor.var import TensorVariable
 from pandas import Series
@@ -532,6 +533,13 @@ class Model(Factor, WithMemoization, metaclass=ContextMeta):
         parameters can only take on valid values you can set this to
         False for increased speed. This should not be used if your model
         contains discrete variables.
+    rng_seeder: int or numpy.random.RandomState
+        The ``numpy.random.RandomState`` used to seed the
+        ``RandomStateSharedVariable`` sequence used by a model
+        ``RandomVariable``s, or an int used to seed a new
+        ``numpy.random.RandomState``.  If ``None``, a
+        ``RandomStateSharedVariable`` will be generated and used.  Incremental
+        access to the state sequence is provided by ``Model.next_rng``.
 
     Examples
     --------
@@ -615,17 +623,31 @@ def __new__(cls, *args, **kwargs):
         instance._aesara_config = kwargs.get("aesara_config", {})
         return instance
 
-    def __init__(self, name="", model=None, aesara_config=None, coords=None, check_bounds=True):
+    def __init__(
+        self,
+        name="",
+        model=None,
+        aesara_config=None,
+        coords=None,
+        check_bounds=True,
+        rng_seeder: Optional[Union[int, np.random.RandomState]] = None,
+    ):
         self.name = name
         self._coords = {}
         self._RV_dims = {}
         self._dim_lengths = {}
         self.add_coords(coords)
         self.check_bounds = check_bounds
 
-        self.default_rng = aesara.shared(np.random.RandomState(), name="default_rng", borrow=True)
-        self.default_rng.tag.is_rng = True
-        self.default_rng.default_update = self.default_rng
+        if rng_seeder is None:
+            self.rng_seeder = np.random.RandomState()
+        elif isinstance(rng_seeder, int):
+            self.rng_seeder = np.random.RandomState(rng_seeder)
+        else:
+            self.rng_seeder = rng_seeder
+
+        # The sequence of model-generated RNGs
+        self.rng_seq = []
 
         if self.parent is not None:
             self.named_vars = treedict(parent=self.parent.named_vars)
@@ -931,6 +953,20 @@ def cont_vars(self):
         """All the continuous variables in the model"""
         return list(typefilter(self.value_vars, continuous_types))
 
+    def next_rng(self) -> RandomStateSharedVariable:
+        """Generate a new ``RandomStateSharedVariable``.
+
+        The new ``RandomStateSharedVariable`` is also added to
+        ``Model.rng_seq``.
+        """
+        new_seed = self.rng_seeder.randint(2 ** 30, dtype=np.int64)
+        next_rng = aesara.shared(np.random.RandomState(new_seed), borrow=True)
+        next_rng.tag.is_rng = True
+
+        self.rng_seq.append(next_rng)
+
+        return next_rng
+
     def shape_from_dims(self, dims):
         shape = []
         if len(set(dims)) != len(dims):

pymc3/sampling.py

@@ -339,7 +339,8 @@ def sample(
         time until completion ("expected time of arrival"; ETA).
     model : Model (optional if in ``with`` context)
     random_seed : int or list of ints
-        A list is accepted if ``cores`` is greater than one.
+        Random seed(s) used by the sampling steps.  A list is accepted if
+        ``cores`` is greater than one.
     discard_tuned_samples : bool
         Whether to discard posterior samples of the tune interval.
     compute_convergence_checks : bool, default=True
@@ -467,10 +468,6 @@ def sample(
             np.random.seed(random_seed)
         random_seed = [np.random.randint(2 ** 30) for _ in range(chains)]
 
-    # TODO: We need to do something about multiple seeds and this single,
-    # shared RNG state.
-    model.default_rng.get_value(borrow=True).seed(random_seed)
-
     if not isinstance(random_seed, abc.Iterable):
         raise TypeError("Invalid value for `random_seed`. Must be tuple, list or int")
 
@@ -1004,9 +1001,7 @@ def _iter_sample(
     """
     model = modelcontext(model)
     draws = int(draws)
-    if random_seed is not None:
-        # np.random.seed(random_seed)
-        model.default_rng.get_value(borrow=True).seed(random_seed)
+
     if draws < 1:
         raise ValueError("Argument `draws` must be greater than 0.")
 
@@ -1273,9 +1268,7 @@ def _prepare_iter_population(
     nchains = len(chains)
     model = modelcontext(model)
     draws = int(draws)
-    if random_seed is not None:
-        # np.random.seed(random_seed)
-        model.default_rng.get_value(borrow=True).seed(random_seed)
+
     if draws < 1:
         raise ValueError("Argument `draws` should be above 0.")
 
@@ -1693,8 +1686,12 @@ def sample_posterior_predictive(
         vars_ = model.observed_RVs + model.auto_deterministics
 
     if random_seed is not None:
-        # np.random.seed(random_seed)
-        model.default_rng.get_value(borrow=True).seed(random_seed)
+        warnings.warn(
+            "In this version, RNG seeding is managed by the Model objects. "
+            "See the `rng_seeder` argument in Model's constructor.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
 
     indices = np.arange(samples)
 
@@ -1816,8 +1813,6 @@ def sample_posterior_predictive_w(
         Dictionary with the variables as keys. The values corresponding to the
         posterior predictive samples from the weighted models.
     """
-    # np.random.seed(random_seed)
-
     if isinstance(traces[0], InferenceData):
         n_samples = [
             trace.posterior.sizes["chain"] * trace.posterior.sizes["draw"] for trace in traces
@@ -1832,9 +1827,15 @@ def sample_posterior_predictive_w(
     if models is None:
         models = [modelcontext(models)] * len(traces)
 
+    if random_seed:
+        warnings.warn(
+            "In this version, RNG seeding is managed by the Model objects. "
+            "See the `rng_seeder` argument in Model's constructor.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     for model in models:
-        if random_seed:
-            model.default_rng.get_value(borrow=True).seed(random_seed)
         if model.potentials:
             warnings.warn(
                 "The effect of Potentials on other parameters is ignored during posterior predictive sampling. "
@@ -1937,6 +1938,7 @@ def sample_prior_predictive(
     model: Optional[Model] = None,
     var_names: Optional[Iterable[str]] = None,
     random_seed=None,
+    mode: Optional[Union[str, Mode]] = None,
 ) -> Dict[str, np.ndarray]:
     """Generate samples from the prior predictive distribution.
 
@@ -1950,6 +1952,8 @@ def sample_prior_predictive(
         samples. Defaults to both observed and unobserved RVs.
     random_seed : int
         Seed for the random number generator.
+    mode:
+        The mode used by ``aesara.function`` to compile the graph.
 
     Returns
     -------
@@ -1977,8 +1981,12 @@ def sample_prior_predictive(
         vars_ = set(var_names)
 
     if random_seed is not None:
-        # np.random.seed(random_seed)
-        model.default_rng.get_value(borrow=True).seed(random_seed)
+        warnings.warn(
+            "In this version, RNG seeding is managed by the Model objects. "
+            "See the `rng_seeder` argument in Model's constructor.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
 
     names = get_default_varnames(vars_, include_transformed=False)
 
@@ -2127,8 +2135,6 @@ def init_nuts(
 
     if random_seed is not None:
         random_seed = int(np.atleast_1d(random_seed)[0])
-        # np.random.seed(random_seed)
-        model.default_rng.get_value(borrow=True).seed(random_seed)
 
     cb = [
         pm.callbacks.CheckParametersConvergence(tolerance=1e-2, diff="absolute"),

pymc3/tests/test_distributions.py

@@ -2944,3 +2944,25 @@ def func(x):
     import pickle
 
     pickle.loads(pickle.dumps(y))
+
+
+def test_distinct_rvs():
+    """Make sure `RandomVariable`s generated using a `Model`'s default RNG state all have distinct states."""
+
+    with pm.Model(rng_seeder=np.random.RandomState(2023532)) as model:
+        X_rv = pm.Normal("x")
+        Y_rv = pm.Normal("y")
+
+        pp_samples = pm.sample_prior_predictive(samples=2)
+
+    assert X_rv.owner.inputs[0] != Y_rv.owner.inputs[0]
+
+    assert len(model.rng_seq) == 2
+
+    with pm.Model(rng_seeder=np.random.RandomState(2023532)):
+        X_rv = pm.Normal("x")
+        Y_rv = pm.Normal("y")
+
+        pp_samples_2 = pm.sample_prior_predictive(samples=2)
+
+    assert np.array_equal(pp_samples["y"], pp_samples_2["y"])

pymc3/tests/test_ndarray_backend.py

@@ -209,8 +209,8 @@ def test_combine_true_squeeze_true(self):
 
 class TestSaveLoad:
     @staticmethod
-    def model():
-        with pm.Model() as model:
+    def model(rng_seeder=None):
+        with pm.Model(rng_seeder=rng_seeder) as model:
             x = pm.Normal("x", 0, 1)
             y = pm.Normal("y", x, 1, observed=2)
             z = pm.Normal("z", x + y, 1)
@@ -267,15 +267,16 @@ def test_sample_posterior_predictive(self, tmpdir_factory):
 
         assert save_dir == directory
 
-        with TestSaveLoad.model() as model:
-            model.default_rng.get_value(borrow=True).seed(10)
+        rng = np.random.RandomState(10)
+
+        with TestSaveLoad.model(rng_seeder=rng):
             ppc = pm.sample_posterior_predictive(self.trace)
 
-        with TestSaveLoad.model() as model:
+        rng = np.random.RandomState(10)
+
+        with TestSaveLoad.model(rng_seeder=rng):
             trace2 = pm.load_trace(directory)
-            model.default_rng.get_value(borrow=True).seed(10)
             ppc2 = pm.sample_posterior_predictive(trace2)
-            ppc2f = pm.sample_posterior_predictive(trace2)
 
         for key, value in ppc.items():
             assert (value == ppc2[key]).all()

pymc3/tests/test_sampling.py

@@ -619,11 +619,13 @@ def test_model_not_drawable_prior(self):
             assert samples["foo"].shape == (40, 200)
 
     def test_model_shared_variable(self):
-        x = np.random.randn(100)
+        rng = np.random.RandomState(9832)
+
+        x = rng.randn(100)
         y = x > 0
         x_shared = aesara.shared(x)
         y_shared = aesara.shared(y)
-        with pm.Model() as model:
+        with pm.Model(rng_seeder=rng) as model:
             coeff = pm.Normal("x", mu=0, sd=1)
             logistic = pm.Deterministic("p", pm.math.sigmoid(coeff * x_shared))
 
@@ -644,12 +646,12 @@ def test_model_shared_variable(self):
         npt.assert_allclose(post_pred["p"], expected_p)
 
     def test_deterministic_of_observed(self):
-        np.random.seed(8442)
+        rng = np.random.RandomState(8442)
 
-        meas_in_1 = pm.aesaraf.floatX(2 + 4 * np.random.randn(10))
-        meas_in_2 = pm.aesaraf.floatX(5 + 4 * np.random.randn(10))
+        meas_in_1 = pm.aesaraf.floatX(2 + 4 * rng.randn(10))
+        meas_in_2 = pm.aesaraf.floatX(5 + 4 * rng.randn(10))
         nchains = 2
-        with pm.Model() as model:
+        with pm.Model(rng_seeder=rng) as model:
             mu_in_1 = pm.Normal("mu_in_1", 0, 1)
             sigma_in_1 = pm.HalfNormal("sd_in_1", 1)
             mu_in_2 = pm.Normal("mu_in_2", 0, 1)
@@ -660,7 +662,6 @@ def test_deterministic_of_observed(self):
             out_diff = in_1 + in_2
             pm.Deterministic("out", out_diff)
 
-            model.default_rng.get_value(borrow=True).seed(0)
             trace = pm.sample(
                 100,
                 chains=nchains,
@@ -670,7 +671,6 @@ def test_deterministic_of_observed(self):
 
             rtol = 1e-5 if aesara.config.floatX == "float64" else 1e-4
 
-            model.default_rng.get_value(borrow=True).seed(0)
             ppc = pm.sample_posterior_predictive(
                 model=model,
                 trace=trace,
@@ -682,11 +682,11 @@ def test_deterministic_of_observed(self):
             npt.assert_allclose(ppc["in_1"] + ppc["in_2"], ppc["out"], rtol=rtol)
 
     def test_deterministic_of_observed_modified_interface(self):
-        np.random.seed(4982)
+        rng = np.random.RandomState(4982)
 
-        meas_in_1 = pm.aesaraf.floatX(2 + 4 * np.random.randn(100))
-        meas_in_2 = pm.aesaraf.floatX(5 + 4 * np.random.randn(100))
-        with pm.Model() as model:
+        meas_in_1 = pm.aesaraf.floatX(2 + 4 * rng.randn(100))
+        meas_in_2 = pm.aesaraf.floatX(5 + 4 * rng.randn(100))
+        with pm.Model(rng_seeder=rng) as model:
             mu_in_1 = pm.Normal("mu_in_1", 0, 1, testval=0)
             sigma_in_1 = pm.HalfNormal("sd_in_1", 1, testval=1)
             mu_in_2 = pm.Normal("mu_in_2", 0, 1, testval=0)
@@ -969,12 +969,10 @@ def test_multivariate2(self):
         assert sim_ppc["obs"].shape == (20,) + mn_data.shape
 
     def test_layers(self):
-        with pm.Model() as model:
+        with pm.Model(rng_seeder=232093) as model:
             a = pm.Uniform("a", lower=0, upper=1, size=10)
             b = pm.Binomial("b", n=1, p=a, size=10)
 
-        model.default_rng.get_value(borrow=True).seed(232093)
-
         b_sampler = aesara.function([], b)
         avg = np.stack([b_sampler() for i in range(10000)]).mean(0)
         npt.assert_array_almost_equal(avg, 0.5 * np.ones((10,)), decimal=2)

pymc3/tests/test_step.py

@@ -1657,7 +1657,9 @@ def perform(self, node, inputs, outputs):
                 mout = []
                 coarse_models = []
 
-                with Model() as coarse_model_0:
+                rng = np.random.RandomState(seed)
+
+                with Model(rng_seeder=rng) as coarse_model_0:
                     if aesara.config.floatX == "float32":
                         Q = Data("Q", np.float32(0.0))
                     else:
@@ -1674,9 +1676,9 @@ def perform(self, node, inputs, outputs):
 
                     coarse_models.append(coarse_model_0)
 
-                coarse_model_0.default_rng.get_value(borrow=True).seed(seed)
+                rng = np.random.RandomState(seed)
 
-                with Model() as coarse_model_1:
+                with Model(rng_seeder=rng) as coarse_model_1:
                     if aesara.config.floatX == "float32":
                         Q = Data("Q", np.float32(0.0))
                     else:
@@ -1693,9 +1695,9 @@ def perform(self, node, inputs, outputs):
 
                     coarse_models.append(coarse_model_1)
 
-                coarse_model_1.default_rng.get_value(borrow=True).seed(seed)
+                rng = np.random.RandomState(seed)
 
-                with Model() as model:
+                with Model(rng_seeder=rng) as model:
                     if aesara.config.floatX == "float32":
                         Q = Data("Q", np.float32(0.0))
                     else: