Refactor Beta to use custom rng_fn clipped_beta_rv

Author: ricardoV94

pymc3/distributions/continuous.py

@@ -23,7 +23,7 @@
 
 from aesara.assert_op import Assert
 from aesara.tensor.random.basic import (
-    beta,
+    BetaRV,
     cauchy,
     exponential,
     gamma,
@@ -42,6 +42,7 @@
     SplineWrapper,
     betaln,
     bound,
+    clipped_beta_rvs,
     gammaln,
     i0e,
     incomplete_beta,
@@ -1075,6 +1076,15 @@ def logcdf(self, value):
         )
 
 
+class BetaClippedRV(BetaRV):
+    @classmethod
+    def rng_fn(cls, rng, alpha, beta, size):
+        return clipped_beta_rvs(alpha, beta, size=size, random_state=rng)
+
+
+beta = BetaClippedRV()
+
+
 class Beta(UnitContinuous):
     r"""
     Beta log-likelihood.
@@ -1149,9 +1159,6 @@ def dist(cls, alpha=None, beta=None, mu=None, sigma=None, sd=None, *args, **kwar
         alpha = aet.as_tensor_variable(floatX(alpha))
         beta = aet.as_tensor_variable(floatX(beta))
 
-        # mean = alpha / (alpha + beta)
-        # variance = (alpha * beta) / ((alpha + beta) ** 2 * (alpha + beta + 1))
-
         assert_negative_support(alpha, "alpha", "Beta")
         assert_negative_support(beta, "beta", "Beta")
 

pymc3/distributions/dist_math.py

@@ -600,7 +600,7 @@ def incomplete_beta(a, b, value):
     )
 
 
-def clipped_beta_rvs(a, b, size=None, dtype="float64"):
+def clipped_beta_rvs(a, b, size=None, random_state=None, dtype="float64"):
     """Draw beta distributed random samples in the open :math:`(0, 1)` interval.
 
     The samples are generated with ``scipy.stats.beta.rvs``, but any value that
@@ -635,6 +635,6 @@ def clipped_beta_rvs(a, b, size=None, dtype="float64"):
         is shifted to ``np.nextafter(1, 0, dtype=dtype)``.
 
     """
-    out = scipy.stats.beta.rvs(a, b, size=size).astype(dtype)
+    out = scipy.stats.beta.rvs(a, b, size=size, random_state=random_state).astype(dtype)
     lower, upper = _beta_clip_values[dtype]
     return np.maximum(np.minimum(out, upper), lower)

pymc3/tests/test_distributions_random.py

@@ -301,7 +301,6 @@ class TestWald(BaseTestCases.BaseTestCase):
     params = {"mu": 1.0, "lam": 1.0, "alpha": 0.0}
 
 
-@pytest.mark.skip(reason="This test is covered by Aesara")
 class TestBeta(BaseTestCases.BaseTestCase):
     distribution = pm.Beta
     params = {"alpha": 1.0, "beta": 1.0}