Repository: aesara-devs/aesara

Issue:
LogSumExp of `-inf` returns `nan` instead of `-inf`
```python
x = at.vector('x')
res1 = at.log(at.sum(at.exp(x)))  # nan with sum
res2 = at.log(at.prod(at.exp(x)))  # but not with prod

fun = aesara.function([x], [res1, res2])
print(fun(np.array([-np.inf, -np.inf])))  # [array(nan), array(-inf)]
```

for reference in numpy we get `-inf`, together with a divide by zero encountered in log

```python
print(np.log(np.sum(np.exp([-np.inf, -np.inf]))))  # -inf
```

This showed up in the pymc3 `logsumexp` function which returns `nan` with this input:

```python
from pymc3.math import logsumexp

x = at.vector('x')
res = logsumexp(x)
fun = aesara.function([x], res) 

print(fun(np.array([-np.inf, -np.inf])))  # [nan]
```
Whereas the scipy reference works fine
```python
print(scipy.special.logsumexp([-np.inf, -np.inf]))  # -inf
```

Weirdly it happens with addition, but not subtraction or multiplication, so maybe the problem is not the sum but the addition?

```python
x = at.vector('x')
res1 = at.log(at.exp(x[0]) + at.exp(x[1]))
res2 = at.log(at.exp(x[0]) - at.exp(x[1]))
fun = aesara.function([x], [res1, res2])
print(fun(np.array([-np.inf, -np.inf])))  # [array(nan), array(-inf)]

Comments:

Patch:
diff --git a/aesara/tensor/math_opt.py b/aesara/tensor/math_opt.py
index e90d6b425..ad0eeb92b 100644
--- a/aesara/tensor/math_opt.py
+++ b/aesara/tensor/math_opt.py
@@ -11,7 +11,6 @@ from functools import reduce
 import numpy as np
 
 import aesara.scalar.basic as aes
-from aesara import compile
 from aesara.assert_op import assert_op
 from aesara.configdefaults import config
 from aesara.graph.basic import Constant, Variable
@@ -74,6 +73,7 @@ from aesara.tensor.math import (
     expm1,
     ge,
     int_div,
+    isinf,
     log,
     log1p,
     makeKeepDims,
@@ -2286,34 +2286,33 @@ def local_log1p(fgraph, node):
 @register_stabilize
 @register_specialize
 @local_optimizer([log])
-def local_log_add(fgraph, node):
-    # log(exp(x)+exp(y))
-    #
-    # Suppose x >= y
-    # log(exp(x) + exp(y))
-    # log(exp(x) * (1 + exp(y)/exp(x)))
-    # x + log(1 + exp(y)/exp(x))
-    # x + log1p(exp(y)/exp(x))
-    # x + log1p(exp(y-x))
+def local_log_add_exp(fgraph, node):
+    # log(exp(x)+exp(y)+exp(z)) = max + log(x-max, y-max, z-max)
+
     if node.op == log:
         z = node.inputs[0]
         if z.owner and z.owner.op == add:
             zi = z.owner.inputs
-            if len(zi) != 2:
-                # -- upgrading Maximum to handle multiple inputs wasn't trivial
-                #    TODO
-                # raise NotImplementedError()
-                return
             pre_exp = [x.owner.inputs[0] for x in zi if x.owner and x.owner.op == exp]
+            # all arguments to add are exp(<something>)
             if len(pre_exp) == len(zi):
-                # all arguments to add are exp(<something>)
-                max_pre = maximum(*pre_exp)
-
-                ret = max_pre + log1p(exp(add(*[p - max_pre for p in pre_exp])))
-                ret.tag.values_eq_approx = values_eq_approx_remove_inf
+                # Do not offset when max_pre = -np.inf, to avoid nan in the output
+                # Switch statement is placed directly inside add to break the self-symmetry
+                # of the returned output (otherwise the optimization would not stabilize)
+                max_pre = reduce(maximum, pre_exp)
+                ret = max_pre + log(
+                    add(
+                        *[
+                            switch(isinf(max_pre), exp(max_pre), exp(p - max_pre))
+                            for p in pre_exp
+                        ]
+                    )
+                )
                 return [ret]
 
 
+@register_stabilize
+@register_specialize
 @local_optimizer([log])
 def local_log_sum_exp(fgraph, node):
     # log(sum_i(exp(x_i))) = x_max + log(sum_i(exp(x_i - x_max)))
@@ -2342,7 +2341,19 @@ def local_log_sum_exp(fgraph, node):
     max_pre_exp = aet_max(pre_exp, axis=axis)
     max_pre_exp_keepdims = makeKeepDims(pre_exp, max_pre_exp, axis)
 
-    ret = max_pre_exp + log(aet_sum(exp(pre_exp - max_pre_exp_keepdims), axis=axis))
+    # Do not offset when max_pre = -np.inf, to avoid nan in the output
+    # Switch statement is placed directly inside sum to break the self-symmetry
+    # of the returned output (otherwise the optimization would not stabilize)
+    ret = max_pre_exp + log(
+        aet_sum(
+            switch(
+                isinf(max_pre_exp_keepdims),
+                exp(max_pre_exp_keepdims),
+                exp(pre_exp - max_pre_exp_keepdims),
+            ),
+            axis=axis,
+        ),
+    )
 
     # Restore the dimshuffle op, if any.
     if dimshuffle_op:
@@ -2351,14 +2362,6 @@ def local_log_sum_exp(fgraph, node):
     return [ret]
 
 
-compile.optdb.register(
-    "local_log_sum_exp",
-    in2out(local_log_sum_exp, ignore_newtrees=True),
-    1.6,
-    "fast_run",
-)
-
-
 def add_calculate(num, denum, aslist=False, out_type=None):
     # TODO: make sure that this function and mul_calculate are similar
     if out_type is None:


Test patch:
diff --git a/tests/tensor/test_math_opt.py b/tests/tensor/test_math_opt.py
index 98a1d8a80..55dfbae0b 100644
--- a/tests/tensor/test_math_opt.py
+++ b/tests/tensor/test_math_opt.py
@@ -1840,10 +1840,7 @@ def test_log1p():
     assert [node.op for node in f.maker.fgraph.toposort()] == [log1p]
 
 
-@pytest.mark.xfail(
-    reason="log(add(exp)) is not stabilized when adding more than 2 elements, see #623"
-)
-def test_log_add():
+def test_local_log_add_exp():
     m = config.mode
     if m == "FAST_COMPILE":
         m = "FAST_RUN"
@@ -1858,26 +1855,28 @@ def test_log_add():
     y = dvector()
     f = function([x, y], log(exp(x) + exp(y)), mode=m)
 
-    f([10000], [10000])  # causes overflow if handled incorrectly
-    assert np.isfinite(f([10000], [10000]))
+    # test that it gives the correct result when it doesn't overflow
+    f([10], [10])  # doesn't causes overflow
+    utt.assert_allclose(f([10], [10]), 10 + np.log1p(1))
+
+    assert np.isfinite(f([10000], [10000]))  # causes overflow if handled incorrectly
     utt.assert_allclose(f([10000], [10000]), 10000 + np.log1p(1))
 
-    # test that it give the same result when it don't overflow
-    f([10], [10])  # don't causes overflow
-    utt.assert_allclose(f([10], [10]), 10 + np.log1p(1))
+    # test that when max = +-inf, optimized output still works correctly
+    assert f([-np.inf], [-np.inf]) == -np.inf
+    assert f([np.inf], [np.inf]) == np.inf
+    assert f([np.inf], [-np.inf]) == np.inf
 
-    # test that it also works with more than two args, (this currently fails)
+    # test that it also works with more than two args
     x = dvector()
     y = dvector()
     f = function([x, y], log(exp(x) + exp(y) + exp(x - y) + exp(x + y)), mode=m)
 
-    f([10000], [10000])  # causes overflow if handled incorrectly
+    assert np.isfinite(f([10000], [10000]))  # causes overflow if handled incorrectly
     utt.assert_allclose(f([10000], [10000]), 20000)
 
     # TODO: test that the optimization works in the presence of broadcasting.
 
-    # TODO: (write and) test that the optimization works with Sum in addition to working with Add.
-
 
 def test_local_subtensor_of_dot():
     m1 = matrix()
@@ -4001,6 +4000,16 @@ def test_local_log_sum_exp3():
     assert np.allclose(optimised_ret, 100.0)
 
 
+def test_local_log_sum_exp_inf():
+    # Test that when max = +-inf, optimized output still works correctly
+    x = vector("x")
+    f = compile_graph_log_sum_exp(x, axis=0)
+
+    assert f([-np.inf, -np.inf]) == -np.inf
+    assert f([np.inf, np.inf]) == np.inf
+    assert f([-np.inf, np.inf]) == np.inf
+
+
 def test_local_reciprocal_1_plus_exp():
     x = vector("x")
     y = aet.reciprocal(1 + exp(x))