From 3b51d3a44232de1b87bd9497fcc2856bcf79ba19 Mon Sep 17 00:00:00 2001
From: Lucca Zenobio <luccazen@gmail.com>
Date: Tue, 8 Jun 2021 21:53:41 -0300
Subject: [PATCH 01/12] refactor special dist

---
 pymc3/distributions/continuous.py     |  2 +-
 pymc3/distributions/dist_math.py      | 40 +++++++++++++++++-
 pymc3/distributions/multivariate.py   |  3 +-
 pymc3/distributions/special.py        | 58 ---------------------------
 pymc3/tests/test_dist_math.py         | 25 ++++++++++++
 pymc3/tests/test_special_functions.py | 45 ---------------------
 6 files changed, 66 insertions(+), 107 deletions(-)
 delete mode 100644 pymc3/distributions/special.py
 delete mode 100644 pymc3/tests/test_special_functions.py

diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py
index 193c6453b4..c3a6543b49 100644
--- a/pymc3/distributions/continuous.py
+++ b/pymc3/distributions/continuous.py
@@ -61,12 +61,12 @@
     incomplete_beta,
     log_normal,
     logpow,
+    log_i0,
     normal_lccdf,
     normal_lcdf,
     zvalue,
 )
 from pymc3.distributions.distribution import Continuous
-from pymc3.distributions.special import log_i0
 from pymc3.math import log1mexp, log1pexp, logdiffexp, logit
 from pymc3.util import UNSET
 
diff --git a/pymc3/distributions/dist_math.py b/pymc3/distributions/dist_math.py
index 7f4112d745..328ba24baf 100644
--- a/pymc3/distributions/dist_math.py
+++ b/pymc3/distributions/dist_math.py
@@ -30,11 +30,11 @@
 from aesara.scalar import UnaryScalarOp, upgrade_to_float_no_complex
 from aesara.scan import until
 from aesara.tensor.elemwise import Elemwise
+from aesara.tensor.math import gammaln, psi
 from aesara.tensor.slinalg import Cholesky, Solve
 
 from pymc3.aesaraf import floatX
 from pymc3.distributions.shape_utils import to_tuple
-from pymc3.distributions.special import gammaln
 
 f = floatX
 c = -0.5 * np.log(2.0 * np.pi)
@@ -634,3 +634,41 @@ def clipped_beta_rvs(a, b, size=None, random_state=None, dtype="float64"):
     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)
+
+
+def multigammaln(a, p):
+    """Multivariate Log Gamma
+
+    Parameters
+    ----------
+    a: tensor like
+    p: int
+       degrees of freedom. p > 0
+    """
+    i = at.arange(1, p + 1)
+    return p * (p - 1) * at.log(np.pi) / 4.0 + at.sum(gammaln(a + (1.0 - i) / 2.0), axis=0)
+
+
+def log_i0(x):
+    """
+    Calculates the logarithm of the 0 order modified Bessel function of the first kind""
+    """
+    return at.switch(
+        at.lt(x, 5),
+        at.log1p(
+            x ** 2.0 / 4.0
+            + x ** 4.0 / 64.0
+            + x ** 6.0 / 2304.0
+            + x ** 8.0 / 147456.0
+            + x ** 10.0 / 14745600.0
+            + x ** 12.0 / 2123366400.0
+        ),
+        x
+        - 0.5 * at.log(2.0 * np.pi * x)
+        + at.log1p(
+            1.0 / (8.0 * x)
+            + 9.0 / (128.0 * x ** 2.0)
+            + 225.0 / (3072.0 * x ** 3.0)
+            + 11025.0 / (98304.0 * x ** 4.0)
+        ),
+    )
diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 7cea4a90e2..8d67bd6e49 100644
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -41,9 +41,8 @@
 from pymc3.aesaraf import floatX, intX
 from pymc3.distributions import transforms
 from pymc3.distributions.continuous import ChiSquared, Normal
-from pymc3.distributions.dist_math import bound, factln, logpow
+from pymc3.distributions.dist_math import bound, factln, logpow, gammaln, multigammaln
 from pymc3.distributions.distribution import Continuous, Discrete
-from pymc3.distributions.special import gammaln, multigammaln
 from pymc3.math import kron_diag, kron_dot, kron_solve_lower, kronecker
 
 __all__ = [
diff --git a/pymc3/distributions/special.py b/pymc3/distributions/special.py
deleted file mode 100644
index 7c0a15e8b2..0000000000
--- a/pymc3/distributions/special.py
+++ /dev/null
@@ -1,58 +0,0 @@
-#   Copyright 2020 The PyMC Developers
-#
-#   Licensed under the Apache License, Version 2.0 (the "License");
-#   you may not use this file except in compliance with the License.
-#   You may obtain a copy of the License at
-#
-#       http://www.apache.org/licenses/LICENSE-2.0
-#
-#   Unless required by applicable law or agreed to in writing, software
-#   distributed under the License is distributed on an "AS IS" BASIS,
-#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#   See the License for the specific language governing permissions and
-#   limitations under the License.
-
-import aesara.tensor as at
-import numpy as np
-
-from aesara.tensor.math import gammaln, psi
-
-__all__ = ["gammaln", "multigammaln", "psi", "log_i0"]
-
-
-def multigammaln(a, p):
-    """Multivariate Log Gamma
-
-    Parameters
-    ----------
-    a: tensor like
-    p: int
-       degrees of freedom. p > 0
-    """
-    i = at.arange(1, p + 1)
-    return p * (p - 1) * at.log(np.pi) / 4.0 + at.sum(gammaln(a + (1.0 - i) / 2.0), axis=0)
-
-
-def log_i0(x):
-    """
-    Calculates the logarithm of the 0 order modified Bessel function of the first kind""
-    """
-    return at.switch(
-        at.lt(x, 5),
-        at.log1p(
-            x ** 2.0 / 4.0
-            + x ** 4.0 / 64.0
-            + x ** 6.0 / 2304.0
-            + x ** 8.0 / 147456.0
-            + x ** 10.0 / 14745600.0
-            + x ** 12.0 / 2123366400.0
-        ),
-        x
-        - 0.5 * at.log(2.0 * np.pi * x)
-        + at.log1p(
-            1.0 / (8.0 * x)
-            + 9.0 / (128.0 * x ** 2.0)
-            + 225.0 / (3072.0 * x ** 3.0)
-            + 11025.0 / (98304.0 * x ** 4.0)
-        ),
-    )
diff --git a/pymc3/tests/test_dist_math.py b/pymc3/tests/test_dist_math.py
index fc1e531a00..7f139baca3 100644
--- a/pymc3/tests/test_dist_math.py
+++ b/pymc3/tests/test_dist_math.py
@@ -17,8 +17,10 @@
 import numpy.testing as npt
 import pytest
 
+from aesara import config, function
 from aesara.tensor.random.basic import multinomial
 from scipy import interpolate, stats
+import scipy.special as ss
 
 import pymc3 as pm
 
@@ -32,8 +34,10 @@
     clipped_beta_rvs,
     factln,
     i0e,
+    multigammaln
 )
 from pymc3.tests.helpers import verify_grad
+from pymc3.tests.checks import close_to
 
 
 def test_bound():
@@ -236,3 +240,24 @@ def test_clipped_beta_rvs(dtype):
     # equal to zero or one (issue #3898)
     values = clipped_beta_rvs(0.01, 0.01, size=1000000, dtype=dtype)
     assert not (np.any(values == 0) or np.any(values == 1))
+
+def check_vals(fn1, fn2, *args):
+    v = fn1(*args)
+    close_to(v, fn2(*args), 1e-6 if v.dtype == np.float64 else 1e-4)
+
+
+def test_multigamma():
+    x = at.vector("x")
+    p = at.scalar("p")
+
+    xvals = [np.array([v], dtype=config.floatX) for v in [0.1, 2, 5, 10, 50, 100]]
+
+    multigammaln_ = function([x, p], multigammaln(x, p), mode="FAST_COMPILE")
+
+    def ssmultigammaln(a, b):
+        return np.array(ss.multigammaln(a[0], b), config.floatX)
+
+    for p in [0, 1, 2, 3, 4, 100]:
+        for x in xvals:
+            if np.all(x > 0.5 * (p - 1)):
+                check_vals(multigammaln_, ssmultigammaln, x, p)
diff --git a/pymc3/tests/test_special_functions.py b/pymc3/tests/test_special_functions.py
deleted file mode 100644
index 9a73647fd1..0000000000
--- a/pymc3/tests/test_special_functions.py
+++ /dev/null
@@ -1,45 +0,0 @@
-#   Copyright 2020 The PyMC Developers
-#
-#   Licensed under the Apache License, Version 2.0 (the "License");
-#   you may not use this file except in compliance with the License.
-#   You may obtain a copy of the License at
-#
-#       http://www.apache.org/licenses/LICENSE-2.0
-#
-#   Unless required by applicable law or agreed to in writing, software
-#   distributed under the License is distributed on an "AS IS" BASIS,
-#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#   See the License for the specific language governing permissions and
-#   limitations under the License.
-
-import aesara.tensor as at
-import numpy as np
-import scipy.special as ss
-
-from aesara import config, function
-
-import pymc3.distributions.special as ps
-
-from pymc3.tests.checks import close_to
-
-
-def check_vals(fn1, fn2, *args):
-    v = fn1(*args)
-    close_to(v, fn2(*args), 1e-6 if v.dtype == np.float64 else 1e-4)
-
-
-def test_multigamma():
-    x = at.vector("x")
-    p = at.scalar("p")
-
-    xvals = [np.array([v], dtype=config.floatX) for v in [0.1, 2, 5, 10, 50, 100]]
-
-    multigammaln = function([x, p], ps.multigammaln(x, p), mode="FAST_COMPILE")
-
-    def ssmultigammaln(a, b):
-        return np.array(ss.multigammaln(a[0], b), config.floatX)
-
-    for p in [0, 1, 2, 3, 4, 100]:
-        for x in xvals:
-            if np.all(x > 0.5 * (p - 1)):
-                check_vals(multigammaln, ssmultigammaln, x, p)

From 6b9548260fedf44a19ed4ac46c7bfc06a7d4c910 Mon Sep 17 00:00:00 2001
From: Lucca Zenobio <luccazen@gmail.com>
Date: Tue, 8 Jun 2021 22:00:53 -0300
Subject: [PATCH 02/12] linter

---
 pymc3/distributions/continuous.py   | 2 +-
 pymc3/distributions/dist_math.py    | 2 +-
 pymc3/distributions/multivariate.py | 2 +-
 pymc3/tests/test_dist_math.py       | 9 ++++++---
 4 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py
index c3a6543b49..c2cff511c9 100644
--- a/pymc3/distributions/continuous.py
+++ b/pymc3/distributions/continuous.py
@@ -59,9 +59,9 @@
     gammaln,
     i0e,
     incomplete_beta,
+    log_i0,
     log_normal,
     logpow,
-    log_i0,
     normal_lccdf,
     normal_lcdf,
     zvalue,
diff --git a/pymc3/distributions/dist_math.py b/pymc3/distributions/dist_math.py
index 328ba24baf..c91117feb8 100644
--- a/pymc3/distributions/dist_math.py
+++ b/pymc3/distributions/dist_math.py
@@ -30,7 +30,7 @@
 from aesara.scalar import UnaryScalarOp, upgrade_to_float_no_complex
 from aesara.scan import until
 from aesara.tensor.elemwise import Elemwise
-from aesara.tensor.math import gammaln, psi
+from aesara.tensor.math import gammaln
 from aesara.tensor.slinalg import Cholesky, Solve
 
 from pymc3.aesaraf import floatX
diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 8d67bd6e49..e07ea95ef8 100644
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -41,7 +41,7 @@
 from pymc3.aesaraf import floatX, intX
 from pymc3.distributions import transforms
 from pymc3.distributions.continuous import ChiSquared, Normal
-from pymc3.distributions.dist_math import bound, factln, logpow, gammaln, multigammaln
+from pymc3.distributions.dist_math import bound, factln, gammaln, logpow, multigammaln
 from pymc3.distributions.distribution import Continuous, Discrete
 from pymc3.math import kron_diag, kron_dot, kron_solve_lower, kronecker
 
diff --git a/pymc3/tests/test_dist_math.py b/pymc3/tests/test_dist_math.py
index 7f139baca3..ccc49178a0 100644
--- a/pymc3/tests/test_dist_math.py
+++ b/pymc3/tests/test_dist_math.py
@@ -16,11 +16,12 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
+import scipy.special as ss
 
 from aesara import config, function
 from aesara.tensor.random.basic import multinomial
 from scipy import interpolate, stats
-import scipy.special as ss
+
 
 import pymc3 as pm
 
@@ -34,10 +35,12 @@
     clipped_beta_rvs,
     factln,
     i0e,
-    multigammaln
+    multigammaln,
 )
-from pymc3.tests.helpers import verify_grad
+
 from pymc3.tests.checks import close_to
+from pymc3.tests.helpers import verify_grad
+
 
 
 def test_bound():

From 21ce563c22cd65c0efe7fadb21c34704a00ea118 Mon Sep 17 00:00:00 2001
From: Lucca Zenobio <luccazen@gmail.com>
Date: Tue, 8 Jun 2021 22:02:25 -0300
Subject: [PATCH 03/12] linter

---
 pymc3/tests/test_dist_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc3/tests/test_dist_math.py b/pymc3/tests/test_dist_math.py
index ccc49178a0..43fd857051 100644
--- a/pymc3/tests/test_dist_math.py
+++ b/pymc3/tests/test_dist_math.py
@@ -42,7 +42,6 @@
 from pymc3.tests.helpers import verify_grad
 
 
-
 def test_bound():
     logp = at.ones((10, 10))
     cond = at.ones((10, 10))
@@ -244,6 +243,7 @@ def test_clipped_beta_rvs(dtype):
     values = clipped_beta_rvs(0.01, 0.01, size=1000000, dtype=dtype)
     assert not (np.any(values == 0) or np.any(values == 1))
 
+
 def check_vals(fn1, fn2, *args):
     v = fn1(*args)
     close_to(v, fn2(*args), 1e-6 if v.dtype == np.float64 else 1e-4)

From 8f4226ee719fefb41821464399c90e9b47c333c2 Mon Sep 17 00:00:00 2001
From: Lucca Zenobio <luccazen@gmail.com>
Date: Tue, 8 Jun 2021 22:04:41 -0300
Subject: [PATCH 04/12] fix pytest

---
 .github/workflows/pytest.yml | 2 --
 1 file changed, 2 deletions(-)

diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index 8a1aa3611f..7bdd280bd7 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -47,7 +47,6 @@ jobs:
             --ignore=pymc3/tests/test_minibatches.py
             --ignore=pymc3/tests/test_pickling.py
             --ignore=pymc3/tests/test_plots.py
-            --ignore=pymc3/tests/test_special_functions.py
             --ignore=pymc3/tests/test_updates.py
             --ignore=pymc3/tests/test_examples.py
             --ignore=pymc3/tests/test_gp.py
@@ -67,7 +66,6 @@ jobs:
             pymc3/tests/test_minibatches.py
             pymc3/tests/test_pickling.py
             pymc3/tests/test_plots.py
-            pymc3/tests/test_special_functions.py
             pymc3/tests/test_updates.py
 
           - |

From 2e4472ad79d753d0fdb740f0efe79002ee2ddfe7 Mon Sep 17 00:00:00 2001
From: Lucca Zenobio <luccazen@gmail.com>
Date: Tue, 8 Jun 2021 22:08:55 -0300
Subject: [PATCH 05/12] fix pytest

---
 pymc3/tests/test_dist_math.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pymc3/tests/test_dist_math.py b/pymc3/tests/test_dist_math.py
index 43fd857051..48ac49ea20 100644
--- a/pymc3/tests/test_dist_math.py
+++ b/pymc3/tests/test_dist_math.py
@@ -22,7 +22,6 @@
 from aesara.tensor.random.basic import multinomial
 from scipy import interpolate, stats
 
-
 import pymc3 as pm
 
 from pymc3.aesaraf import floatX
@@ -37,7 +36,6 @@
     i0e,
     multigammaln,
 )
-
 from pymc3.tests.checks import close_to
 from pymc3.tests.helpers import verify_grad
 

From 09100caad6982b9a1439666af23ac0729d0fb698 Mon Sep 17 00:00:00 2001
From: Lucca Zenobio <luccazen@gmail.com>
Date: Wed, 9 Jun 2021 19:14:50 -0300
Subject: [PATCH 06/12] refactor tests

---
 pymc3/tests/test_distributions.py | 4 ++++
 pymc3/tests/test_ode.py           | 6 +++---
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/pymc3/tests/test_distributions.py b/pymc3/tests/test_distributions.py
index 12f08eab32..c0e6380755 100644
--- a/pymc3/tests/test_distributions.py
+++ b/pymc3/tests/test_distributions.py
@@ -1356,6 +1356,8 @@ def test_t(self):
             lambda value, nu, mu, lam: sp.t.logcdf(value, nu, mu, lam ** -0.5),
             n_samples=10,  # relies on slow incomplete beta
         )
+        # TODO: reenable when PR #4736 is merged
+        """
         self.check_logcdf(
             StudentT,
             R,
@@ -1363,6 +1365,8 @@ def test_t(self):
             lambda value, nu, mu, sigma: sp.t.logcdf(value, nu, mu, sigma),
             n_samples=5,  # Just testing alternative parametrization
         )
+        """
+        
 
     def test_cauchy(self):
         self.check_logp(
diff --git a/pymc3/tests/test_ode.py b/pymc3/tests/test_ode.py
index 95937a3674..5b8b878d8b 100644
--- a/pymc3/tests/test_ode.py
+++ b/pymc3/tests/test_ode.py
@@ -395,9 +395,9 @@ def system(y, t, p):
         ode_model = DifferentialEquation(func=system, t0=0, times=times, n_states=2, n_theta=2)
 
         with pm.Model() as model:
-            beta = pm.HalfCauchy("beta", 1)
-            gamma = pm.HalfCauchy("gamma", 1)
-            sigma = pm.HalfCauchy("sigma", 1, shape=2)
+            beta = pm.HalfCauchy("beta", 1, initval=1)
+            gamma = pm.HalfCauchy("gamma", 1, initval=1)
+            sigma = pm.HalfCauchy("sigma", 1, shape=2, initval=[1, 1])
             forward = ode_model(theta=[beta, gamma], y0=[0.99, 0.01])
             y = pm.Lognormal("y", mu=pm.math.log(forward), sd=sigma, observed=yobs)
 

From 7f325a5d94d589b89b6a953a5a605bc084c34b7a Mon Sep 17 00:00:00 2001
From: Lucca Zenobio <luccazen@gmail.com>
Date: Wed, 9 Jun 2021 19:16:18 -0300
Subject: [PATCH 07/12] refactor tests

---
 pymc3/tests/test_distributions.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pymc3/tests/test_distributions.py b/pymc3/tests/test_distributions.py
index c0e6380755..33644253e4 100644
--- a/pymc3/tests/test_distributions.py
+++ b/pymc3/tests/test_distributions.py
@@ -1366,7 +1366,6 @@ def test_t(self):
             n_samples=5,  # Just testing alternative parametrization
         )
         """
-        
 
     def test_cauchy(self):
         self.check_logp(

From 32e9e28aaf838b2624f10065d3032ae48f3d26bc Mon Sep 17 00:00:00 2001
From: Lucca Zenobio <luccazen@gmail.com>
Date: Thu, 10 Jun 2021 10:32:22 -0300
Subject: [PATCH 08/12] fixes

---
 pymc3/distributions/continuous.py   |  2 +-
 pymc3/distributions/dist_math.py    | 11 +++++------
 pymc3/distributions/multivariate.py | 27 ++++++++++++++-------------
 pymc3/tests/test_dist_math.py       |  4 ++--
 4 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py
index c2cff511c9..c0dbe9fc53 100644
--- a/pymc3/distributions/continuous.py
+++ b/pymc3/distributions/continuous.py
@@ -23,6 +23,7 @@
 import aesara.tensor as at
 import numpy as np
 
+from aesara.tensor import gammaln
 from aesara.assert_op import Assert
 from aesara.tensor.random.basic import (
     BetaRV,
@@ -56,7 +57,6 @@
     betaln,
     bound,
     clipped_beta_rvs,
-    gammaln,
     i0e,
     incomplete_beta,
     log_i0,
diff --git a/pymc3/distributions/dist_math.py b/pymc3/distributions/dist_math.py
index c91117feb8..81e22c492f 100644
--- a/pymc3/distributions/dist_math.py
+++ b/pymc3/distributions/dist_math.py
@@ -30,7 +30,6 @@
 from aesara.scalar import UnaryScalarOp, upgrade_to_float_no_complex
 from aesara.scan import until
 from aesara.tensor.elemwise import Elemwise
-from aesara.tensor.math import gammaln
 from aesara.tensor.slinalg import Cholesky, Solve
 
 from pymc3.aesaraf import floatX
@@ -110,7 +109,7 @@ def logpow(x, m):
 
 
 def factln(n):
-    return gammaln(n + 1)
+    return at.gammaln(n + 1)
 
 
 def binomln(n, k):
@@ -118,7 +117,7 @@ def binomln(n, k):
 
 
 def betaln(x, y):
-    return gammaln(x) + gammaln(y) - gammaln(x + y)
+    return at.gammaln(x) + at.gammaln(y) - at.gammaln(x + y)
 
 
 def std_cdf(x):
@@ -553,7 +552,7 @@ def _step(i, t, s, a, b, value):
 
     s = s[-1] + t1 + ai
 
-    t = gammaln(a + b) - gammaln(a) - gammaln(b) + a * at.log(value) + at.log(s)
+    t = at.gammaln(a + b) - at.gammaln(a) - at.gammaln(b) + a * at.log(value) + at.log(s)
     return at.exp(t)
 
 
@@ -585,7 +584,7 @@ def incomplete_beta(a, b, value):
 
     # Direct incomplete beta accounting for flipped a, b.
     t = at.exp(
-        a * at.log(x) + b * at.log(xc) + gammaln(a + b) - gammaln(a) - gammaln(b) + at.log(w / a)
+        a * at.log(x) + b * at.log(xc) + at.gammaln(a + b) - at.gammaln(a) - at.gammaln(b) + at.log(w / a)
     )
 
     t = at.switch(flip, at.switch(at.le(t, machep), one - machep, one - t), t)
@@ -646,7 +645,7 @@ def multigammaln(a, p):
        degrees of freedom. p > 0
     """
     i = at.arange(1, p + 1)
-    return p * (p - 1) * at.log(np.pi) / 4.0 + at.sum(gammaln(a + (1.0 - i) / 2.0), axis=0)
+    return p * (p - 1) * at.log(np.pi) / 4.0 + at.sum(at.gammaln(a + (1.0 - i) / 2.0), axis=0)
 
 
 def log_i0(x):
diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index e07ea95ef8..94e2387352 100644
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -22,6 +22,7 @@
 import numpy as np
 import scipy
 
+from aesara.tensor import gammaln
 from aesara.graph.basic import Apply
 from aesara.graph.op import Op
 from aesara.tensor.nlinalg import det, eigh, matrix_inverse, trace
@@ -41,7 +42,7 @@
 from pymc3.aesaraf import floatX, intX
 from pymc3.distributions import transforms
 from pymc3.distributions.continuous import ChiSquared, Normal
-from pymc3.distributions.dist_math import bound, factln, gammaln, logpow, multigammaln
+from pymc3.distributions.dist_math import bound, factln, logpow, multigammaln
 from pymc3.distributions.distribution import Continuous, Discrete
 from pymc3.math import kron_diag, kron_dot, kron_solve_lower, kronecker
 
@@ -153,7 +154,7 @@ def quaddist_tau(delta, chol_mat):
 
 
 class MvNormal(Continuous):
-    R"""
+    r"""
     Multivariate normal log-likelihood.
 
     .. math::
@@ -248,7 +249,7 @@ def _distr_parameters_for_repr(self):
 
 
 class MvStudentT(Continuous):
-    R"""
+    r"""
     Multivariate Student-T log-likelihood.
 
     .. math::
@@ -361,7 +362,7 @@ def _distr_parameters_for_repr(self):
 
 
 class Dirichlet(Continuous):
-    R"""
+    r"""
     Dirichlet log-likelihood.
 
     .. math::
@@ -451,7 +452,7 @@ def rng_fn(cls, rng, n, p, size):
 
 
 class Multinomial(Discrete):
-    R"""
+    r"""
     Multinomial log-likelihood.
 
     Generalizes binomial distribution, but instead of each trial resulting
@@ -524,7 +525,7 @@ def logp(value, n, p):
 
 
 class DirichletMultinomial(Discrete):
-    R"""Dirichlet Multinomial log-likelihood.
+    r"""Dirichlet Multinomial log-likelihood.
 
     Dirichlet mixture of Multinomials distribution, with a marginalized PMF.
 
@@ -728,7 +729,7 @@ def __str__(self):
 
 
 class Wishart(Continuous):
-    R"""
+    r"""
     Wishart log-likelihood.
 
     The Wishart distribution is the probability distribution of the
@@ -945,7 +946,7 @@ def _lkj_normalizing_constant(eta, n):
 
 
 class _LKJCholeskyCov(Continuous):
-    R"""Underlying class for covariance matrix with LKJ distributed correlations.
+    r"""Underlying class for covariance matrix with LKJ distributed correlations.
     See docs for LKJCholeskyCov function for more details on how to use it in models.
     """
 
@@ -1125,7 +1126,7 @@ def _distr_parameters_for_repr(self):
 
 
 def LKJCholeskyCov(name, eta, n, sd_dist, compute_corr=False, store_in_trace=True, *args, **kwargs):
-    R"""Wrapper function for covariance matrix with LKJ distributed correlations.
+    r"""Wrapper function for covariance matrix with LKJ distributed correlations.
 
     This defines a distribution over Cholesky decomposed covariance
     matrices, such that the underlying correlation matrices follow an
@@ -1278,7 +1279,7 @@ def LKJCholeskyCov(name, eta, n, sd_dist, compute_corr=False, store_in_trace=Tru
 
 
 class LKJCorr(Continuous):
-    R"""
+    r"""
     The LKJ (Lewandowski, Kurowicka and Joe) log-likelihood.
 
     The LKJ distribution is a prior distribution for correlation matrices.
@@ -1434,7 +1435,7 @@ def _distr_parameters_for_repr(self):
 
 
 class MatrixNormal(Continuous):
-    R"""
+    r"""
     Matrix-valued normal log-likelihood.
 
     .. math::
@@ -1693,7 +1694,7 @@ def _distr_parameters_for_repr(self):
 
 
 class KroneckerNormal(Continuous):
-    R"""
+    r"""
     Multivariate normal log-likelihood with Kronecker-structured covariance.
 
     .. math::
@@ -1940,7 +1941,7 @@ def _distr_parameters_for_repr(self):
 
 
 class CAR(Continuous):
-    R"""
+    r"""
     Likelihood for a conditional autoregression. This is a special case of the
     multivariate normal with an adjacency-structured covariance matrix.
 
diff --git a/pymc3/tests/test_dist_math.py b/pymc3/tests/test_dist_math.py
index 48ac49ea20..4223ba4c3d 100644
--- a/pymc3/tests/test_dist_math.py
+++ b/pymc3/tests/test_dist_math.py
@@ -16,7 +16,7 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
-import scipy.special as ss
+import scipy.special as st
 
 from aesara import config, function
 from aesara.tensor.random.basic import multinomial
@@ -256,7 +256,7 @@ def test_multigamma():
     multigammaln_ = function([x, p], multigammaln(x, p), mode="FAST_COMPILE")
 
     def ssmultigammaln(a, b):
-        return np.array(ss.multigammaln(a[0], b), config.floatX)
+        return np.array(st.multigammaln(a[0], b), config.floatX)
 
     for p in [0, 1, 2, 3, 4, 100]:
         for x in xvals:

From 50b4ea8106df9973fdfd4795207835897ca8c637 Mon Sep 17 00:00:00 2001
From: Lucca Zenobio <luccazen@gmail.com>
Date: Thu, 10 Jun 2021 10:36:35 -0300
Subject: [PATCH 09/12] fixes

---
 pymc3/distributions/continuous.py   |  2 +-
 pymc3/distributions/dist_math.py    | 11 ++++++-----
 pymc3/distributions/multivariate.py |  2 +-
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py
index c0dbe9fc53..d5893b5e24 100644
--- a/pymc3/distributions/continuous.py
+++ b/pymc3/distributions/continuous.py
@@ -23,8 +23,8 @@
 import aesara.tensor as at
 import numpy as np
 
-from aesara.tensor import gammaln
 from aesara.assert_op import Assert
+from aesara.tensor import gammaln
 from aesara.tensor.random.basic import (
     BetaRV,
     WeibullRV,
diff --git a/pymc3/distributions/dist_math.py b/pymc3/distributions/dist_math.py
index 81e22c492f..2935d41a1f 100644
--- a/pymc3/distributions/dist_math.py
+++ b/pymc3/distributions/dist_math.py
@@ -27,6 +27,7 @@
 from aesara.compile.builders import OpFromGraph
 from aesara.graph.basic import Apply
 from aesara.graph.op import Op
+from aesara.tensor import gammaln
 from aesara.scalar import UnaryScalarOp, upgrade_to_float_no_complex
 from aesara.scan import until
 from aesara.tensor.elemwise import Elemwise
@@ -109,7 +110,7 @@ def logpow(x, m):
 
 
 def factln(n):
-    return at.gammaln(n + 1)
+    return gammaln(n + 1)
 
 
 def binomln(n, k):
@@ -117,7 +118,7 @@ def binomln(n, k):
 
 
 def betaln(x, y):
-    return at.gammaln(x) + at.gammaln(y) - at.gammaln(x + y)
+    return gammaln(x) + gammaln(y) - gammaln(x + y)
 
 
 def std_cdf(x):
@@ -552,7 +553,7 @@ def _step(i, t, s, a, b, value):
 
     s = s[-1] + t1 + ai
 
-    t = at.gammaln(a + b) - at.gammaln(a) - at.gammaln(b) + a * at.log(value) + at.log(s)
+    t = gammaln(a + b) - gammaln(a) - gammaln(b) + a * at.log(value) + at.log(s)
     return at.exp(t)
 
 
@@ -584,7 +585,7 @@ def incomplete_beta(a, b, value):
 
     # Direct incomplete beta accounting for flipped a, b.
     t = at.exp(
-        a * at.log(x) + b * at.log(xc) + at.gammaln(a + b) - at.gammaln(a) - at.gammaln(b) + at.log(w / a)
+        a * at.log(x) + b * at.log(xc) + gammaln(a + b) - gammaln(a) - gammaln(b) + at.log(w / a)
     )
 
     t = at.switch(flip, at.switch(at.le(t, machep), one - machep, one - t), t)
@@ -645,7 +646,7 @@ def multigammaln(a, p):
        degrees of freedom. p > 0
     """
     i = at.arange(1, p + 1)
-    return p * (p - 1) * at.log(np.pi) / 4.0 + at.sum(at.gammaln(a + (1.0 - i) / 2.0), axis=0)
+    return p * (p - 1) * at.log(np.pi) / 4.0 + at.sum(gammaln(a + (1.0 - i) / 2.0), axis=0)
 
 
 def log_i0(x):
diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 94e2387352..758f1492e6 100644
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -22,9 +22,9 @@
 import numpy as np
 import scipy
 
-from aesara.tensor import gammaln
 from aesara.graph.basic import Apply
 from aesara.graph.op import Op
+from aesara.tensor import gammaln
 from aesara.tensor.nlinalg import det, eigh, matrix_inverse, trace
 from aesara.tensor.random.basic import MultinomialRV, dirichlet, multivariate_normal
 from aesara.tensor.random.utils import broadcast_params

From 193e934e546a07d9f4b7fb1b6f77a1f779e8f528 Mon Sep 17 00:00:00 2001
From: Lucca Zenobio <luccazen@gmail.com>
Date: Thu, 10 Jun 2021 10:38:24 -0300
Subject: [PATCH 10/12] fixes

---
 pymc3/distributions/dist_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc3/distributions/dist_math.py b/pymc3/distributions/dist_math.py
index 2935d41a1f..491a81c32e 100644
--- a/pymc3/distributions/dist_math.py
+++ b/pymc3/distributions/dist_math.py
@@ -27,9 +27,9 @@
 from aesara.compile.builders import OpFromGraph
 from aesara.graph.basic import Apply
 from aesara.graph.op import Op
-from aesara.tensor import gammaln
 from aesara.scalar import UnaryScalarOp, upgrade_to_float_no_complex
 from aesara.scan import until
+from aesara.tensor import gammaln
 from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.slinalg import Cholesky, Solve
 

From 72ee28d2670f1f663a667e391efc480f158ed6ba Mon Sep 17 00:00:00 2001
From: Ricardo <ricardo.vieira1994@gmail.com>
Date: Thu, 10 Jun 2021 19:16:58 +0200
Subject: [PATCH 11/12] small renaming

---
 pymc3/tests/test_dist_math.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pymc3/tests/test_dist_math.py b/pymc3/tests/test_dist_math.py
index 4223ba4c3d..ee53671727 100644
--- a/pymc3/tests/test_dist_math.py
+++ b/pymc3/tests/test_dist_math.py
@@ -16,7 +16,7 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
-import scipy.special as st
+import scipy.special as sp
 
 from aesara import config, function
 from aesara.tensor.random.basic import multinomial
@@ -255,10 +255,10 @@ def test_multigamma():
 
     multigammaln_ = function([x, p], multigammaln(x, p), mode="FAST_COMPILE")
 
-    def ssmultigammaln(a, b):
-        return np.array(st.multigammaln(a[0], b), config.floatX)
+    def ref_multigammaln(a, b):
+        return np.array(sp.multigammaln(a[0], b), config.floatX)
 
     for p in [0, 1, 2, 3, 4, 100]:
         for x in xvals:
             if np.all(x > 0.5 * (p - 1)):
-                check_vals(multigammaln_, ssmultigammaln, x, p)
+                check_vals(multigammaln_, ref_multigammaln, x, p)

From ccba84d6b61fe269acd514e752db1803066e74df Mon Sep 17 00:00:00 2001
From: Ricardo <ricardo.vieira1994@gmail.com>
Date: Thu, 10 Jun 2021 19:23:20 +0200
Subject: [PATCH 12/12] explicit import

---
 pymc3/tests/test_dist_math.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pymc3/tests/test_dist_math.py b/pymc3/tests/test_dist_math.py
index ee53671727..c9389990ed 100644
--- a/pymc3/tests/test_dist_math.py
+++ b/pymc3/tests/test_dist_math.py
@@ -16,7 +16,7 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
-import scipy.special as sp
+import scipy.special
 
 from aesara import config, function
 from aesara.tensor.random.basic import multinomial
@@ -256,7 +256,7 @@ def test_multigamma():
     multigammaln_ = function([x, p], multigammaln(x, p), mode="FAST_COMPILE")
 
     def ref_multigammaln(a, b):
-        return np.array(sp.multigammaln(a[0], b), config.floatX)
+        return np.array(scipy.special.multigammaln(a[0], b), config.floatX)
 
     for p in [0, 1, 2, 3, 4, 100]:
         for x in xvals: