Add log1pexp_numpy helper

pymc/math.py

@@ -356,6 +356,25 @@ def log1mexp_numpy(x, *, negative_input=False):
     return out
 
 
+def log1pexp_numpy(x):
+    """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
+    """
+    return np.where(
+        x < -37.0,
+        np.exp(x),
+        np.where(
+            x < 33.3,
+            np.log1p(np.exp(x)),
+            x,
+        ),
+    )
+
+
 def flatten_list(tensors):
     return at.concatenate([var.ravel() for var in tensors])
 

pymc/tests/test_math.py

@@ -18,6 +18,8 @@
 import numpy.testing as npt
 import pytest
 
+from aesara.tensor import log1pexp
+
 from pymc.aesaraf import floatX
 from pymc.math import (
     LogDet,
@@ -30,6 +32,7 @@
     kronecker,
     log1mexp,
     log1mexp_numpy,
+    log1pexp_numpy,
     log_softmax,
     logdet,
     logdiffexp,
@@ -188,6 +191,16 @@ def test_log1mexp_deprecation_warnings():
     assert np.isclose(res_neg_at, res_neg)
 
 
+def test_log1pexp_numpy():
+    a = np.array([-40, -30, 0, 30, 40])
+    assert np.all(
+        np.isclose(
+            log1pexp_numpy(a),
+            log1pexp(a).eval(),
+        )
+    )
+
+
 def test_logdiffexp():
     a = np.log([1, 2, 3, 4])
     b = np.log([0, 1, 2, 3])