Remove custom logaddexp, logsumexp, and log1pexp functions and deprecate log1mexp in favor of Aesara implementations.

Closes #4747

Author: ricardoV94

RELEASE-NOTES.md

@@ -30,6 +30,7 @@
 - The length of `dims` in the model is now tracked symbolically through `Model.dim_lengths` (see [#4625](https://github.com/pymc-devs/pymc3/pull/4625)).
 - We now include `cloudpickle` as a required dependency, and no longer depend on `dill` (see [#4858](https://github.com/pymc-devs/pymc3/pull/4858)).
 - The `incomplete_beta` function in `pymc3.distributions.dist_math` was replaced by `aesara.tensor.betainc` (see [4857](https://github.com/pymc-devs/pymc3/pull/4857)).
+- `math.log1mexp` and `math.log1mexp_numpy` will expect negative inputs in the future. A `FutureWarning` is now raised unless `negative_input=True` is set (see [#4860](https://github.com/pymc-devs/pymc3/pull/4860)).
 - ...
 
 ## PyMC3 3.11.2 (14 March 2021)

pymc3/distributions/continuous.py

@@ -67,7 +67,7 @@
     zvalue,
 )
 from pymc3.distributions.distribution import Continuous
-from pymc3.math import log1mexp, log1pexp, logdiffexp, logit
+from pymc3.math import logdiffexp, logit
 from pymc3.util import UNSET
 
 __all__ = [
@@ -1095,7 +1095,7 @@ def logcdf(
         return bound(
             at.switch(
                 at.lt(value, np.inf),
-                a + log1pexp(b - a),
+                a + at.log1pexp(b - a),
                 0,
             ),
             0 < value,
@@ -1370,7 +1370,7 @@ def logcdf(value, a, b):
         -------
         TensorVariable
         """
-        logcdf = log1mexp(-(b * at.log1p(-(value ** a))))
+        logcdf = at.log1mexp(b * at.log1p(-(value ** a)))
         return bound(at.switch(value < 1, logcdf, 0), value >= 0, a > 0, b > 0)
 
 
@@ -1462,7 +1462,7 @@ def logcdf(value, mu):
         """
         lam = at.inv(mu)
         return bound(
-            log1mexp(lam * value),
+            at.log1mexp(-lam * value),
             0 <= value,
             0 <= lam,
         )
@@ -2711,7 +2711,7 @@ def logcdf(value, alpha, beta):
         """
         a = (value / beta) ** alpha
         return bound(
-            log1mexp(a),
+            at.log1mexp(-a),
             0 <= value,
             0 < alpha,
             0 < beta,
@@ -3662,7 +3662,7 @@ def logcdf(value, mu, s):
         """
 
         return bound(
-            -log1pexp(-(value - mu) / s),
+            -at.log1pexp(-(value - mu) / s),
             0 < s,
         )
 

pymc3/distributions/discrete.py

@@ -42,7 +42,7 @@
 )
 from pymc3.distributions.distribution import Discrete
 from pymc3.distributions.logprob import _logcdf, _logp
-from pymc3.math import log1mexp, logaddexp, logsumexp, sigmoid
+from pymc3.math import sigmoid
 
 __all__ = [
     "Binomial",
@@ -279,7 +279,7 @@ def logcdf(value, n, alpha, beta):
         return bound(
             at.switch(
                 at.lt(value, n),
-                logsumexp(
+                at.logsumexp(
                     BetaBinomial.logp(at.arange(safe_lower, value + 1), n, alpha, beta),
                     keepdims=False,
                 ),
@@ -826,7 +826,7 @@ def logcdf(value, p):
         """
 
         return bound(
-            log1mexp(-at.log1p(-p) * value),
+            at.log1mexp(at.log1p(-p) * value),
             0 <= value,
             0 <= p,
             p <= 1,
@@ -945,7 +945,7 @@ def logcdf(value, good, bad, n):
         return bound(
             at.switch(
                 at.lt(value, n),
-                logsumexp(
+                at.logsumexp(
                     HyperGeometric.logp(at.arange(safe_lower, value + 1), good, bad, n),
                     keepdims=False,
                 ),
@@ -1300,7 +1300,7 @@ def logp(value, psi, theta):
         logp_val = at.switch(
             at.gt(value, 0),
             at.log(psi) + _logp(poisson, value, {}, theta),
-            logaddexp(at.log1p(-psi), at.log(psi) - theta),
+            at.logaddexp(at.log1p(-psi), at.log(psi) - theta),
         )
 
         return bound(
@@ -1328,7 +1328,7 @@ def logcdf(value, psi, theta):
         """
 
         return bound(
-            logaddexp(at.log1p(-psi), at.log(psi) + _logcdf(poisson, value, {}, theta)),
+            at.logaddexp(at.log1p(-psi), at.log(psi) + _logcdf(poisson, value, {}, theta)),
             0 <= value,
             0 <= psi,
             psi <= 1,
@@ -1430,7 +1430,7 @@ def logp(value, psi, n, p):
         logp_val = at.switch(
             at.gt(value, 0),
             at.log(psi) + _logp(binomial, value, {}, n, p),
-            logaddexp(at.log1p(-psi), at.log(psi) + n * at.log1p(-p)),
+            at.logaddexp(at.log1p(-psi), at.log(psi) + n * at.log1p(-p)),
         )
 
         return bound(
@@ -1460,7 +1460,7 @@ def logcdf(value, psi, n, p):
         """
 
         return bound(
-            logaddexp(at.log1p(-psi), at.log(psi) + _logcdf(binomial, value, {}, n, p)),
+            at.logaddexp(at.log1p(-psi), at.log(psi) + _logcdf(binomial, value, {}, n, p)),
             0 <= value,
             value <= n,
             0 <= psi,
@@ -1583,7 +1583,7 @@ def logp(value, psi, n, p):
             at.switch(
                 at.gt(value, 0),
                 at.log(psi) + _logp(nbinom, value, {}, n, p),
-                logaddexp(at.log1p(-psi), at.log(psi) + n * at.log(p)),
+                at.logaddexp(at.log1p(-psi), at.log(psi) + n * at.log(p)),
             ),
             0 <= value,
             0 <= psi,
@@ -1609,7 +1609,7 @@ def logcdf(value, psi, n, p):
         TensorVariable
         """
         return bound(
-            logaddexp(at.log1p(-psi), at.log(psi) + _logcdf(nbinom, value, {}, n, p)),
+            at.logaddexp(at.log1p(-psi), at.log(psi) + _logcdf(nbinom, value, {}, n, p)),
             0 <= value,
             0 <= psi,
             psi <= 1,

pymc3/math.py

@@ -13,6 +13,7 @@
 #   limitations under the License.
 
 import sys
+import warnings
 
 from functools import partial, reduce
 
@@ -50,6 +51,9 @@
     gt,
     le,
     log,
+    log1pexp,
+    logaddexp,
+    logsumexp,
     lt,
     maximum,
     minimum,
@@ -186,27 +190,14 @@ def tround(*args, **kwargs):
     return at.round(*args, **kwargs)
 
 
-def logsumexp(x, axis=None, keepdims=True):
-    # Adapted from https://github.com/Theano/Theano/issues/1563
-    x_max = at.max(x, axis=axis, keepdims=True)
-    x_max = at.switch(at.isinf(x_max), 0, x_max)
-    res = at.log(at.sum(at.exp(x - x_max), axis=axis, keepdims=True)) + x_max
-    return res if keepdims else res.squeeze()
-
-
-def logaddexp(a, b):
-    diff = b - a
-    return at.switch(diff > 0, b + at.log1p(at.exp(-diff)), a + at.log1p(at.exp(diff)))
-
-
 def logdiffexp(a, b):
     """log(exp(a) - exp(b))"""
-    return a + log1mexp(a - b)
+    return a + at.log1mexp(b - a)
 
 
 def logdiffexp_numpy(a, b):
     """log(exp(a) - exp(b))"""
-    return a + log1mexp_numpy(a - b)
+    return a + log1mexp_numpy(b - a, negative_input=True)
 
 
 def invlogit(x, eps=sys.float_info.epsilon):
@@ -224,15 +215,7 @@ def logit(p):
     return at.log(p / (floatX(1) - p))
 
 
-def log1pexp(x):
-    """Return log(1 + exp(x)), also called softplus.
-
-    This function is numerically more stable than the naive approach.
-    """
-    return at.softplus(x)
-
-
-def log1mexp(x):
+def log1mexp(x, *, negative_input=False):
     r"""Return log(1 - exp(-x)).
 
     This function is numerically more stable than the naive approach.
@@ -246,21 +229,40 @@ def log1mexp(x):
        "Accurately computing `\log(1-\exp(- \mid a \mid))` Assessed by the Rmpfr package"
 
     """
-    return at.switch(at.lt(x, 0.6931471805599453), at.log(-at.expm1(-x)), at.log1p(-at.exp(-x)))
+    if not negative_input:
+        warnings.warn(
+            "pymc3.math.log1mexp will expect a negative input in a future "
+            "version of PyMC3.\n To suppress this warning set `negative_input=True`",
+            FutureWarning,
+            stacklevel=2,
+        )
+        x = -x
 
+    return at.log1mexp(x)
 
-def log1mexp_numpy(x):
-    """Return log(1 - exp(-x)).
+
+def log1mexp_numpy(x, *, negative_input=False):
+    """Return log(1 - exp(x)).
     This function is numerically more stable than the naive approach.
     For details, see
     https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf
     """
-    x = np.asarray(x)
+    x = np.asarray(x, dtype="float")
+
+    if not negative_input:
+        warnings.warn(
+            "pymc3.math.log1mexp_numpy will expect a negative input in a future "
+            "version of PyMC3.\n To suppress this warning set `negative_input=True`",
+            FutureWarning,
+            stacklevel=2,
+        )
+        x = -x
+
     out = np.empty_like(x)
-    mask = x < 0.6931471805599453  # log(2)
-    out[mask] = np.log(-np.expm1(-x[mask]))
+    mask = x < -0.6931471805599453  # log(1/2)
+    out[mask] = np.log1p(-np.exp(x[mask]))
     mask = ~mask
-    out[mask] = np.log1p(-np.exp(-x[mask]))
+    out[mask] = np.log(-np.expm1(x[mask]))
     return out
 
 

pymc3/tests/test_distributions.py

@@ -1165,7 +1165,7 @@ def scipy_log_pdf(value, a, b):
             )
 
         def scipy_log_cdf(value, a, b):
-            return pm.math.log1mexp_numpy(-(b * np.log1p(-(value ** a))))
+            return pm.math.log1mexp_numpy(b * np.log1p(-(value ** a)), negative_input=True)
 
         self.check_logp(
             Kumaraswamy,