Do not create temporary SymbolicDistribution just to retrieve number …

…of RNGs needed

Reordered methods for consistency

pymc/distributions/censored.py

@@ -76,6 +76,14 @@ 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):
 
@@ -96,24 +104,12 @@ def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None):
         rv_out.tag.upper = upper
 
         if rngs is not None:
-            rv_out = cls.change_rngs(rv_out, rngs)
+            rv_out = cls._change_rngs(rv_out, rngs)
 
         return rv_out
 
     @classmethod
-    def ndim_supp(cls, *dist_params):
-        return 0
-
-    @classmethod
-    def change_size(cls, rv, new_size, expand=False):
-        dist = rv.tag.dist
-        lower = rv.tag.lower
-        upper = rv.tag.upper
-        new_dist = change_rv_size(dist, new_size, expand=expand)
-        return cls.rv_op(new_dist, lower, upper)
-
-    @classmethod
-    def change_rngs(cls, rv, new_rngs):
+    def _change_rngs(cls, rv, new_rngs):
         (new_rng,) = new_rngs
         dist_node = rv.tag.dist.owner
         lower = rv.tag.lower
@@ -123,8 +119,12 @@ def change_rngs(cls, rv, new_rngs):
         return cls.rv_op(new_dist, lower, upper)
 
     @classmethod
-    def graph_rvs(cls, rv):
-        return (rv.tag.dist,)
+    def change_size(cls, rv, new_size, expand=False):
+        dist = rv.tag.dist
+        lower = rv.tag.lower
+        upper = rv.tag.upper
+        new_dist = change_rv_size(dist, new_size, expand=expand)
+        return cls.rv_op(new_dist, lower, upper)
 
 
 @_moment.register(Clip)

pymc/distributions/distribution.py

@@ -396,18 +396,19 @@ def __new__(
             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
+            distribution can be defined based on an arbitrary base distribution.
         cls.rv_op
             Returns a TensorVariable that represents the symbolic distribution
             parametrized by a default set of parameters and a size and rngs arguments
-        cls.ndim_supp
-            Returns the support of the symbolic distribution, given the default
-            parameters. This may not always be constant, for instance if the symbolic
-            distribution can be defined based on an arbitrary base distribution.
         cls.change_size
             Returns an equivalent symbolic distribution with a different size. This is
             analogous to `pymc.aesaraf.change_rv_size` for `RandomVariable`s.
-        cls.graph_rvs
-            Returns base RVs in a symbolic distribution.
 
         Parameters
         ----------
@@ -465,9 +466,9 @@ def __new__(
             raise TypeError(f"Name needs to be a string but got: {name}")
 
         if rngs is None:
-            # Create a temporary rv to obtain number of rngs needed
-            temp_graph = cls.dist(*args, rngs=None, **kwargs)
-            rngs = [model.next_rng() for _ in cls.graph_rvs(temp_graph)]
+            # 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]
 

pymc/distributions/mixture.py

@@ -205,6 +205,18 @@ 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
@@ -329,11 +341,6 @@ def _resize_components(cls, size, *components):
 
         return [change_rv_size(component, size) for component in components]
 
-    @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 change_size(cls, rv, new_size, expand=False):
         weights = rv.tag.weights
@@ -355,14 +362,6 @@ def change_size(cls, rv, new_size, expand=False):
 
         return cls.rv_op(weights, *components, rngs=rngs, size=None)
 
-    @classmethod
-    def graph_rvs(cls, rv):
-        # We return rv, which is itself a pseudo RandomVariable, that contains a
-        # mix_indexes_ RV in its inner graph. We want super().dist() to generate
-        # (components + 1) rngs for us, and it will do so based on how many elements
-        # we return here
-        return (*rv.tag.components, rv)
-
 
 @_get_measurable_outputs.register(MarginalMixtureRV)
 def _get_measurable_outputs_MarginalMixtureRV(op, node):