Handle edge case of logdiffexp(-inf, -inf)

Author: ricardoV94

pymc/logprob/transforms.py

@@ -123,6 +123,7 @@
     filter_measurable_variables,
     find_negated_var,
 )
+from pymc.math import logdiffexp
 
 
 class Transform(abc.ABC):
@@ -267,7 +268,7 @@ def measurable_transform_logcdf(op: MeasurableTransform, value, *inputs, **kwarg
             logcdf_zero = _logcdf_helper(measurable_input, 0)
             logcdf = pt.switch(
                 pt.lt(backward_value, 0),
-                pt.log(pt.exp(logcdf_zero) - pt.exp(logcdf)),
+                logdiffexp(logcdf_zero, logcdf),
                 pt.logaddexp(logccdf, logcdf_zero),
             )
     else:

pymc/math.py

@@ -282,7 +282,14 @@ def kron_diag(*diags):
 
 def logdiffexp(a, b):
     """Return log(exp(a) - exp(b))."""
-    return a + pt.log1mexp(b - a)
+    return pt.where(
+        # Handle special case where b is -inf
+        # If a == b == -inf, this will return the correct result of -inf
+        # whereas the default else branch would get a nan due to -inf - (-inf)
+        pt.isneginf(b),
+        a,
+        a + pt.log1mexp(b - a),
+    )
 
 
 invlogit = sigmoid

tests/test_math.py

@@ -144,9 +144,13 @@ def test_log1mexp_deprecation_warnings():
 def test_logdiffexp():
     a = np.log([1, 2, 3, 4])
     with warnings.catch_warnings():
-        warnings.filterwarnings("ignore", "divide by zero encountered in log", RuntimeWarning)
         b = np.log([0, 1, 2, 3])
-    assert np.allclose(logdiffexp(a, b).eval(), 0)
+    np.testing.assert_allclose(logdiffexp(a, b).eval(), 0, atol=1e-15)
+
+    np.testing.assert_allclose(
+        logdiffexp(-np.inf, -np.inf).eval(),
+        -np.inf,
+    )
 
 
 class TestLogDet: