Fix different / alternative parametrizations between RandomOps and Logp methods

.github/workflows/pytest.yml

@@ -26,8 +26,6 @@ jobs:
         # 5th block: These tests PASS without a single XFAIL
         # 6th block: These have some XFAILs
           - |
-            --ignore=pymc3/tests/test_distribution_defaults.py
-            --ignore=pymc3/tests/test_distributions_random.py
             --ignore=pymc3/tests/test_distributions_timeseries.py
             --ignore=pymc3/tests/test_missing.py
             --ignore=pymc3/tests/test_mixture.py
@@ -53,7 +51,6 @@ jobs:
             --ignore=pymc3/tests/test_plots.py
             --ignore=pymc3/tests/test_special_functions.py
             --ignore=pymc3/tests/test_updates.py
-            --ignore=pymc3/tests/test_dist_math.py
             --ignore=pymc3/tests/test_examples.py
             --ignore=pymc3/tests/test_glm.py
             --ignore=pymc3/tests/test_gp.py
@@ -66,8 +63,11 @@ jobs:
             --ignore=pymc3/tests/test_posdef_sym.py
             --ignore=pymc3/tests/test_quadpotential.py
             --ignore=pymc3/tests/test_shape_handling.py
+            --ignore=pymc3/tests/test_distributions.py
+            --ignore=pymc3/tests/test_distributions_random.py
 
           - |
+            pymc3/tests/test_modelcontext.py
             pymc3/tests/test_dist_math.py
             pymc3/tests/test_minibatches.py
             pymc3/tests/test_pickling.py
@@ -76,7 +76,8 @@ jobs:
             pymc3/tests/test_updates.py
 
           - |
-            pymc3/tests/test_dist_math.py
+            pymc3/tests/test_distributions.py
+            pymc3/tests/test_distributions_random.py
             pymc3/tests/test_examples.py
             pymc3/tests/test_glm.py
             pymc3/tests/test_gp.py

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,
@@ -786,18 +787,12 @@ def dist(cls, sigma=None, tau=None, sd=None, *args, **kwargs):
 
         tau, sigma = get_tau_sigma(tau=tau, sigma=sigma)
 
-        # sigma = sd = sigma = aet.as_tensor_variable(sigma)
-        # tau = tau = aet.as_tensor_variable(tau)
-
-        # mean = aet.sqrt(2 / (np.pi * tau))
-        # variance = (1.0 - 2 / np.pi) / tau
-
         assert_negative_support(tau, "tau", "HalfNormal")
         assert_negative_support(sigma, "sigma", "HalfNormal")
 
-        return super().dist([sigma, tau], **kwargs)
+        return super().dist([0.0, sigma], **kwargs)
 
-    def logp(value, sigma, tau):
+    def logp(value, loc, sigma):
         """
         Calculate log-probability of HalfNormal distribution at specified value.
 
@@ -811,14 +806,16 @@ def logp(value, sigma, tau):
         -------
         TensorVariable
         """
+        tau, sigma = get_tau_sigma(tau=None, sigma=sigma)
+
         return bound(
-            -0.5 * tau * value ** 2 + 0.5 * aet.log(tau * 2.0 / np.pi),
-            value >= 0,
+            -0.5 * tau * (value - loc) ** 2 + 0.5 * aet.log(tau * 2.0 / np.pi),
+            value >= loc,
             tau > 0,
             sigma > 0,
         )
 
-    def logcdf(value, sigma, tau):
+    def logcdf(value, loc, sigma):
         """
         Compute the log of the cumulative distribution function for HalfNormal distribution
         at the specified value.
@@ -833,10 +830,10 @@ def logcdf(value, sigma, tau):
         -------
         TensorVariable
         """
-        z = zvalue(value, mu=0, sigma=sigma)
+        z = zvalue(value, mu=loc, sigma=sigma)
         return bound(
             aet.log1p(-aet.erfc(z / aet.sqrt(2.0))),
-            0 <= value,
+            loc <= value,
             0 < sigma,
         )
 
@@ -1079,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.
@@ -1153,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")
 
@@ -2259,14 +2262,10 @@ class HalfCauchy(PositiveContinuous):
     @classmethod
     def dist(cls, beta, *args, **kwargs):
         beta = aet.as_tensor_variable(floatX(beta))
-
-        # mode = aet.as_tensor_variable(0)
-        # median = beta
-
         assert_negative_support(beta, "beta", "HalfCauchy")
-        return super().dist([beta], **kwargs)
+        return super().dist([0.0, beta], **kwargs)
 
-    def logp(value, beta, alpha):
+    def logp(value, loc, beta):
         """
         Calculate log-probability of HalfCauchy distribution at specified value.
 
@@ -2281,12 +2280,12 @@ def logp(value, beta, alpha):
         TensorVariable
         """
         return bound(
-            aet.log(2) - aet.log(np.pi) - aet.log(beta) - aet.log1p((value / beta) ** 2),
-            value >= 0,
+            aet.log(2) - aet.log(np.pi) - aet.log(beta) - aet.log1p(((value - loc) / beta) ** 2),
+            value >= loc,
             beta > 0,
         )
 
-    def logcdf(value, beta, alpha):
+    def logcdf(value, loc, beta):
         """
         Compute the log of the cumulative distribution function for HalfCauchy distribution
         at the specified value.
@@ -2302,8 +2301,8 @@ def logcdf(value, beta, alpha):
         TensorVariable
         """
         return bound(
-            aet.log(2 * aet.arctan(value / beta) / np.pi),
-            0 <= value,
+            aet.log(2 * aet.arctan((value - loc) / beta) / np.pi),
+            loc <= value,
             0 < beta,
         )
 

pymc3/distributions/dist_math.py

@@ -276,7 +276,7 @@ def MvNormalLogp():
     n, k = delta.shape
     n, k = f(n), f(k)
     chol_cov = cholesky(cov)
-    diag = aet.nlinalg.diag(chol_cov)
+    diag = aet.diag(chol_cov)
     ok = aet.all(diag > 0)
 
     chol_cov = aet.switch(ok, chol_cov, aet.fill(chol_cov, 1))
@@ -296,7 +296,7 @@ def dlogp(inputs, gradients):
         n, k = delta.shape
 
         chol_cov = cholesky(cov)
-        diag = aet.nlinalg.diag(chol_cov)
+        diag = aet.diag(chol_cov)
         ok = aet.all(diag > 0)
 
         chol_cov = aet.switch(ok, chol_cov, aet.fill(chol_cov, 1))
@@ -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_dist_math.py

@@ -19,6 +19,7 @@
 import numpy.testing as npt
 import pytest
 
+from aesara.tensor.random.basic import multinomial
 from scipy import interpolate, stats
 
 import pymc3 as pm
@@ -91,16 +92,13 @@ def test_alltrue_shape():
 
 
 class MultinomialA(Discrete):
-    def __init__(self, n, p, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    rv_op = multinomial
 
-        self.n = n
-        self.p = p
-
-    def logp(self, value):
-        n = self.n
-        p = self.p
+    @classmethod
+    def dist(cls, n, p, *args, **kwargs):
+        return super().dist([n, p], **kwargs)
 
+    def logp(value, n, p):
         return bound(
             factln(n) - factln(value).sum() + (value * aet.log(p)).sum(),
             value >= 0,
@@ -112,16 +110,13 @@ def logp(self, value):
 
 
 class MultinomialB(Discrete):
-    def __init__(self, n, p, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self.n = n
-        self.p = p
+    rv_op = multinomial
 
-    def logp(self, value):
-        n = self.n
-        p = self.p
+    @classmethod
+    def dist(cls, n, p, *args, **kwargs):
+        return super().dist([n, p], **kwargs)
 
+    def logp(value, n, p):
         return bound(
             factln(n) - factln(value).sum() + (value * aet.log(p)).sum(),
             aet.all(value >= 0),
@@ -132,26 +127,24 @@ def logp(self, value):
         )
 
 
-@pytest.mark.xfail(reason="This test relies on the deprecated Distribution interface")
 def test_multinomial_bound():
 
     x = np.array([1, 5])
     n = x.sum()
 
     with pm.Model() as modelA:
-        p_a = pm.Dirichlet("p", floatX(np.ones(2)), shape=(2,))
+        p_a = pm.Dirichlet("p", floatX(np.ones(2)))
         MultinomialA("x", n, p_a, observed=x)
 
     with pm.Model() as modelB:
-        p_b = pm.Dirichlet("p", floatX(np.ones(2)), shape=(2,))
+        p_b = pm.Dirichlet("p", floatX(np.ones(2)))
         MultinomialB("x", n, p_b, observed=x)
 
     assert np.isclose(
         modelA.logp({"p_stickbreaking__": [0]}), modelB.logp({"p_stickbreaking__": [0]})
     )
 
 
-@pytest.mark.xfail(reason="MvNormal not implemented")
 class TestMvNormalLogp:
     def test_logp(self):
         np.random.seed(42)
@@ -192,11 +185,10 @@ def func(chol_vec, delta):
         delta_val = floatX(np.random.randn(5, 2))
         verify_grad(func, [chol_vec_val, delta_val])
 
-    @pytest.mark.skip(reason="Fix in aesara not released yet: Theano#5908")
     @aesara.config.change_flags(compute_test_value="ignore")
     def test_hessian(self):
         chol_vec = aet.vector("chol_vec")
-        chol_vec.tag.test_value = np.array([0.1, 2, 3])
+        chol_vec.tag.test_value = floatX(np.array([0.1, 2, 3]))
         chol = aet.stack(
             [
                 aet.stack([aet.exp(0.1 * chol_vec[0]), 0]),
@@ -205,9 +197,10 @@ def test_hessian(self):
         )
         cov = aet.dot(chol, chol.T)
         delta = aet.matrix("delta")
-        delta.tag.test_value = np.ones((5, 2))
+        delta.tag.test_value = floatX(np.ones((5, 2)))
         logp = MvNormalLogp()(cov, delta)
         g_cov, g_delta = aet.grad(logp, [cov, delta])
+        # TODO: What's the test?  Something needs to be asserted.
         aet.grad(g_delta.sum() + g_cov.sum(), [delta, cov])