Remove model.rng_seeder and allow compile_pymc to reseed RNG variable…

…s in compiled function

pymc/aesaraf.py

@@ -922,10 +922,25 @@ def reseed_rngs(rngs: Iterable[SharedVariable], seed: Optional[int]) -> None:
 
 
 def compile_pymc(
-    inputs, outputs, mode=None, **kwargs
+    inputs,
+    outputs,
+    random_seed=None,
+    mode=None,
+    **kwargs,
 ) -> Callable[..., Union[np.ndarray, List[np.ndarray]]]:
     """Use ``aesara.function`` with specialized pymc rewrites always enabled.
 
+    Parameters
+    ----------
+    inputs: list of TensorVariables, optional
+        Inputs of the compiled Aesara function
+    outputs: list of TensorVariables, options
+        Outputs of the compiled Aesara function
+    random_seed: int, optional
+        Seed used override any RandomState/Generator variables in the graph
+    mode: optional
+        Aesara mode used to compile the function
+
     Included rewrites
     -----------------
     random_make_inplace
@@ -945,7 +960,6 @@ def compile_pymc(
     """
     # Create an update mapping of RandomVariable's RNG so that it is automatically
     # updated after every function call
-    # TODO: This won't work for variables with InnerGraphs (Scan and OpFromGraph)
     rng_updates = {}
     output_to_list = outputs if isinstance(outputs, (list, tuple)) else [outputs]
     for random_var in (
@@ -959,11 +973,17 @@ def compile_pymc(
             rng = random_var.owner.inputs[0]
             if not hasattr(rng, "default_update"):
                 rng_updates[rng] = random_var.owner.outputs[0]
+            else:
+                rng_updates[rng] = rng.default_update
         else:
             update_fn = getattr(random_var.owner.op, "update", None)
             if update_fn is not None:
                 rng_updates.update(update_fn(random_var.owner))
 
+    # We always reseed random variables as this provides RNGs with no chances of collision
+    if rng_updates:
+        reseed_rngs(rng_updates.keys(), random_seed)
+
     # If called inside a model context, see if check_bounds flag is set to False
     try:
         from pymc.model import modelcontext

pymc/distributions/censored.py

@@ -76,16 +76,12 @@ def dist(cls, dist, lower, upper, **kwargs):
         check_dist_not_registered(dist)
         return super().dist([dist, lower, upper], **kwargs)
 
-    @classmethod
-    def num_rngs(cls, *args, **kwargs):
-        return 1
-
     @classmethod
     def ndim_supp(cls, *dist_params):
         return 0
 
     @classmethod
-    def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None):
+    def rv_op(cls, dist, lower=None, upper=None, size=None):
 
         lower = at.constant(-np.inf) if lower is None else at.as_tensor_variable(lower)
         upper = at.constant(np.inf) if upper is None else at.as_tensor_variable(upper)
@@ -103,21 +99,8 @@ def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None):
         rv_out.tag.lower = lower
         rv_out.tag.upper = upper
 
-        if rngs is not None:
-            rv_out = cls._change_rngs(rv_out, rngs)
-
         return rv_out
 
-    @classmethod
-    def _change_rngs(cls, rv, new_rngs):
-        (new_rng,) = new_rngs
-        dist_node = rv.tag.dist.owner
-        lower = rv.tag.lower
-        upper = rv.tag.upper
-        olg_rng, size, dtype, *dist_params = dist_node.inputs
-        new_dist = dist_node.op.make_node(new_rng, size, dtype, *dist_params).default_output()
-        return cls.rv_op(new_dist, lower, upper)
-
     @classmethod
     def change_size(cls, rv, new_size, expand=False):
         dist = rv.tag.dist

pymc/distributions/distribution.py

@@ -19,7 +19,7 @@
 
 from abc import ABCMeta
 from functools import singledispatch
-from typing import Callable, Iterable, Optional, Sequence, Tuple, Union, cast
+from typing import Callable, Optional, Sequence, Tuple, Union, cast
 
 import aesara
 import numpy as np
@@ -258,13 +258,10 @@ def __new__(
         if not isinstance(name, string_types):
             raise TypeError(f"Name needs to be a string but got: {name}")
 
-        if rng is None:
-            rng = model.next_rng()
-
         # Create the RV and process dims and observed to determine
         # a shape by which the created RV may need to be resized.
         rv_out, dims, observed, resize_shape = _make_rv_and_resize_shape(
-            cls=cls, dims=dims, model=model, observed=observed, args=args, rng=rng, **kwargs
+            cls=cls, dims=dims, model=model, observed=observed, args=args, **kwargs
         )
 
         if resize_shape:
@@ -383,9 +380,6 @@ class SymbolicDistribution:
         to a canonical parametrization. It should call `super().dist()`, passing a
         list with the default parameters as the first and only non keyword argument,
         followed by other keyword arguments like size and rngs, and return the result
-    cls.num_rngs
-        Returns the number of rngs given the same arguments passed by the user when
-        calling the distribution
     cls.ndim_supp
         Returns the support of the symbolic distribution, given the default set of
         parameters. This may not always be constant, for instance if the symbolic
@@ -402,7 +396,6 @@ def __new__(
         cls,
         name: str,
         *args,
-        rngs: Optional[Iterable] = None,
         dims: Optional[Dims] = None,
         initval=None,
         observed=None,
@@ -419,8 +412,6 @@ def __new__(
             A distribution class that inherits from SymbolicDistribution.
         name : str
             Name for the new model variable.
-        rngs : optional
-            Random number generator to use for the RandomVariable(s) in the graph.
         dims : tuple, optional
             A tuple of dimension names known to the model.
         initval : optional
@@ -468,17 +459,10 @@ def __new__(
         if not isinstance(name, string_types):
             raise TypeError(f"Name needs to be a string but got: {name}")
 
-        if rngs is None:
-            # Instead of passing individual RNG variables we could pass a RandomStream
-            # and let the classes create as many RNGs as they need
-            rngs = [model.next_rng() for _ in range(cls.num_rngs(*args, **kwargs))]
-        elif not isinstance(rngs, (list, tuple)):
-            rngs = [rngs]
-
         # Create the RV and process dims and observed to determine
         # a shape by which the created RV may need to be resized.
         rv_out, dims, observed, resize_shape = _make_rv_and_resize_shape(
-            cls=cls, dims=dims, model=model, observed=observed, args=args, rngs=rngs, **kwargs
+            cls=cls, dims=dims, model=model, observed=observed, args=args, **kwargs
         )
 
         if resize_shape:

pymc/distributions/mixture.py

@@ -205,27 +205,15 @@ def dist(cls, w, comp_dists, **kwargs):
         w = at.as_tensor_variable(w)
         return super().dist([w, *comp_dists], **kwargs)
 
-    @classmethod
-    def num_rngs(cls, w, comp_dists, **kwargs):
-        if not isinstance(comp_dists, (tuple, list)):
-            # comp_dists is a single component
-            comp_dists = [comp_dists]
-        return len(comp_dists) + 1
-
     @classmethod
     def ndim_supp(cls, weights, *components):
         # We already checked that all components have the same support dimensionality
         return components[0].owner.op.ndim_supp
 
     @classmethod
-    def rv_op(cls, weights, *components, size=None, rngs=None):
-        # Update rngs if provided
-        if rngs is not None:
-            components = cls._reseed_components(rngs, *components)
-            *_, mix_indexes_rng = rngs
-        else:
-            # Create new rng for the mix_indexes internal RV
-            mix_indexes_rng = aesara.shared(np.random.default_rng())
+    def rv_op(cls, weights, *components, size=None):
+        # Create new rng for the mix_indexes internal RV
+        mix_indexes_rng = aesara.shared(np.random.default_rng())
 
         single_component = len(components) == 1
         ndim_supp = components[0].owner.op.ndim_supp
@@ -317,19 +305,6 @@ def rv_op(cls, weights, *components, size=None, rngs=None):
 
         return mix_out
 
-    @classmethod
-    def _reseed_components(cls, rngs, *components):
-        *components_rngs, mix_indexes_rng = rngs
-        assert len(components) == len(components_rngs)
-        new_components = []
-        for component, component_rng in zip(components, components_rngs):
-            component_node = component.owner
-            old_rng, *inputs = component_node.inputs
-            new_components.append(
-                component_node.op.make_node(component_rng, *inputs).default_output()
-            )
-        return new_components
-
     @classmethod
     def _resize_components(cls, size, *components):
         if len(components) == 1:
@@ -345,7 +320,6 @@ def _resize_components(cls, size, *components):
     def change_size(cls, rv, new_size, expand=False):
         weights = rv.tag.weights
         components = rv.tag.components
-        rngs = [component.owner.inputs[0] for component in components] + [rv.tag.choices_rng]
 
         if expand:
             component = rv.tag.components[0]
@@ -360,7 +334,7 @@ def change_size(cls, rv, new_size, expand=False):
 
         components = cls._resize_components(new_size, *components)
 
-        return cls.rv_op(weights, *components, rngs=rngs, size=None)
+        return cls.rv_op(weights, *components, size=None)
 
 
 @_get_measurable_outputs.register(MarginalMixtureRV)

pymc/distributions/timeseries.py

@@ -494,28 +494,12 @@ def _get_ar_order(cls, rhos: TensorVariable, ar_order: Optional[int], constant:
 
         return ar_order
 
-    @classmethod
-    def num_rngs(cls, *args, **kwargs):
-        return 2
-
     @classmethod
     def ndim_supp(cls, *args):
         return 1
 
     @classmethod
-    def rv_op(cls, rhos, sigma, init_dist, steps, ar_order, constant_term, size=None, rngs=None):
-
-        if rngs is None:
-            rngs = [
-                aesara.shared(np.random.default_rng(seed))
-                for seed in np.random.SeedSequence().spawn(2)
-            ]
-        (init_dist_rng, noise_rng) = rngs
-        # Re-seed init_dist
-        if init_dist.owner.inputs[0] is not init_dist_rng:
-            _, *inputs = init_dist.owner.inputs
-            init_dist = init_dist.owner.op.make_node(init_dist_rng, *inputs).default_output()
-
+    def rv_op(cls, rhos, sigma, init_dist, steps, ar_order, constant_term, size=None):
         # Init dist should have shape (*size, ar_order)
         if size is not None:
             batch_size = size
@@ -543,6 +527,8 @@ def rv_op(cls, rhos, sigma, init_dist, steps, ar_order, constant_term, size=None
             rhos_bcast_shape_ = (*rhos_bcast_shape_[:-1], rhos_bcast_shape_[-1] + 1)
         rhos_bcast_ = at.broadcast_to(rhos_, rhos_bcast_shape_)
 
+        noise_rng = aesara.shared(np.random.default_rng())
+
         def step(*args):
             *prev_xs, reversed_rhos, sigma, rng = args
             if constant_term:
@@ -581,16 +567,12 @@ def change_size(cls, rv, new_size, expand=False):
             old_size = rv.shape[:-1]
             new_size = at.concatenate([new_size, old_size])
 
-        init_dist_rng = rv.owner.inputs[2].owner.inputs[0]
-        noise_rng = rv.owner.inputs[-1]
-
         op = rv.owner.op
         return cls.rv_op(
             *rv.owner.inputs,
             ar_order=op.ar_order,
             constant_term=op.constant_term,
             size=new_size,
-            rngs=(init_dist_rng, noise_rng),
         )
 
 

pymc/model.py

@@ -44,7 +44,6 @@
 from aesara.graph.basic import Constant, Variable, graph_inputs
 from aesara.graph.fg import FunctionGraph
 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 TensorConstant, TensorVariable
 
@@ -445,13 +444,6 @@ class Model(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
     --------
@@ -549,20 +541,10 @@ def __init__(
         name="",
         coords=None,
         check_bounds=True,
-        rng_seeder: Optional[Union[int, np.random.RandomState]] = None,
     ):
         self.name = self._validate_name(name)
         self.check_bounds = check_bounds
 
-        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: List[SharedVariable] = []
         self._initial_values: Dict[TensorVariable, Optional[Union[np.ndarray, Variable, str]]] = {}
 
         if self.parent is not None:
@@ -1016,8 +998,6 @@ def initial_point(self, seed=None) -> Dict[str, np.ndarray]:
         ip : dict
             Maps names of transformed variables to numeric initial values in the transformed space.
         """
-        if seed is None:
-            seed = self.rng_seeder.randint(2**30, dtype=np.int64)
         fn = make_initial_point_fn(model=self, return_transformed=True)
         return Point(fn(seed), model=self)
 
@@ -1038,20 +1018,6 @@ def set_initval(self, rv_var, initval):
 
         self.initial_values[rv_var] = initval
 
-    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):
@@ -1381,14 +1347,11 @@ def make_obs_var(
                 clone=False,
             )
             (observed_rv_var,) = local_subtensor_rv_lift.transform(fgraph, fgraph.outputs[0].owner)
-            # Make a clone of the RV, but change the rng so that observed and missing
-            # are not treated as equivalent nodes by aesara. This would happen if the
-            # size of the masked and unmasked array happened to coincide
+            # Make a clone of the RV, but let it create a new rng so that observed and
+            # missing are not treated as equivalent nodes by aesara. This would happen
+            # if the size of the masked and unmasked array happened to coincide
             _, size, _, *inps = observed_rv_var.owner.inputs
-            rng = self.model.next_rng()
-            observed_rv_var = observed_rv_var.owner.op(
-                *inps, size=size, rng=rng, name=f"{name}_observed"
-            )
+            observed_rv_var = observed_rv_var.owner.op(*inps, size=size, name=f"{name}_observed")
             observed_rv_var.tag.observations = nonmissing_data
 
             self.create_value_var(observed_rv_var, transform=None, value_var=nonmissing_data)