Support the return_sign argument in logsumexp derivatives

autograd/scipy/special.py

@@ -120,10 +120,22 @@ def gammainc_vjp_arg1(ans, a, x):
 logsumexp = primitive(scipy.special.logsumexp)
 
 
-def make_grad_logsumexp(ans, x, axis=None, b=1.0, keepdims=False):
+def make_grad_logsumexp(ans, x, axis=None, b=1.0, keepdims=False, return_sign=False):
+    if return_sign:
+        # scipy returns (log(abs(sum(b * exp(x)))), sign(sum(b * exp(x)))).
+        ans, sign = ans
     shape, dtype = np.shape(x), np.result_type(x)
 
     def vjp(g):
+        if return_sign:
+            # The sign is piecewise constant, so only the log-sum-exp part of
+            # the output has a non-zero derivative, and the sign of the sum
+            # appears in it because ``ans`` is the log of the absolute value of
+            # that sum.
+            g_repeated, _ = repeat_to_match_shape(g[0], shape, dtype, axis, keepdims)
+            ans_repeated, _ = repeat_to_match_shape(ans, shape, dtype, axis, keepdims)
+            sign_repeated, _ = repeat_to_match_shape(sign, shape, dtype, axis, keepdims)
+            return g_repeated * sign_repeated * b * np.exp(x - ans_repeated)
         g_repeated, _ = repeat_to_match_shape(g, shape, dtype, axis, keepdims)
         ans_repeated, _ = repeat_to_match_shape(ans, shape, dtype, axis, keepdims)
         return g_repeated * b * np.exp(x - ans_repeated)
@@ -134,14 +146,29 @@ def vjp(g):
 defvjp(logsumexp, make_grad_logsumexp)
 
 
-def fwd_grad_logsumexp(g, ans, x, axis=None, b=1.0, keepdims=False):
+def fwd_grad_logsumexp(g, ans, x, axis=None, b=1.0, keepdims=False, return_sign=False):
+    if return_sign:
+        ans, sign = ans
+        # The sign of the sum is piecewise constant, so its tangent is zero and
+        # the shape of the output is the shape of the unexpanded sign.
+        sign_of_sum = sign
+    else:
+        sign = 1.0
     if not keepdims:
         if isinstance(axis, int):
-            ans = np.expand_dims(ans, axis)
+            axes = (axis,)
         elif isinstance(axis, tuple):
-            for ax in sorted(axis):
-                ans = np.expand_dims(ans, ax)
-    return np.sum(g * b * np.exp(x - ans), axis=axis, keepdims=keepdims)
+            axes = axis
+        else:
+            axes = ()
+        for ax in sorted(axes):
+            ans = np.expand_dims(ans, ax)
+            if return_sign:
+                sign = np.expand_dims(sign, ax)
+    result = np.sum(g * sign * b * np.exp(x - ans), axis=axis, keepdims=keepdims)
+    if return_sign:
+        return result, np.zeros_like(sign_of_sum)
+    return result
 
 
 defjvp(logsumexp, fwd_grad_logsumexp)

tests/test_scipy.py

@@ -263,6 +263,39 @@ def f(a):
         check_grads(f, modes=["fwd", "rev"])(x)
         check_grads(lambda a: grad(f)(a), modes=["fwd", "rev"])(x)
 
+    def test_logsumexp_return_sign():
+        # scipy returns a (logsumexp, sign) pair, and the sign belongs to the
+        # gradient: d/dx log(abs(sum(b * exp(x)))) = sign * b * exp(x - ans)
+        f = lambda a, **kwargs: special.logsumexp(a, return_sign=True, **kwargs)[0]
+        combo_check(f, [0], modes=["fwd", "rev"])(
+            [R(4)],
+            b=[npo.array([1.0, -1.0, 1.0, -1.0])],
+            axis=[None, 0],
+            keepdims=[True, False],
+        )
+        combo_check(f, [0], modes=["fwd", "rev"])(
+            [R(3, 4)],
+            b=[npo.exp(R(3, 4)) * npo.array([[1.0], [-1.0], [1.0]])],
+            axis=[None, 0, 1],
+            keepdims=[True, False],
+        )
+
+    def test_logsumexp_negative_sum():
+        # The sum is negative, so the gradient is negated with respect to the
+        # sign-blind formula.
+        x = npo.array([0.0, 0.0])
+        b = npo.array([1.0, -3.0])
+        ans, sign = special.logsumexp(x, b=b, return_sign=True)
+        assert sign == -1.0
+        assert npo.allclose(ans, npo.log(2.0))
+
+        g = grad(lambda a: special.logsumexp(a, b=b, return_sign=True)[0])(x)
+        assert npo.allclose(g, b * npo.exp(x) / (b * npo.exp(x)).sum())
+        check_grads(
+            lambda a: special.logsumexp(a, b=b, return_sign=True)[0],
+            modes=["fwd", "rev"],
+        )(x)
+
     ### Signal ###
     def test_convolve_generalization():
         ag_convolve = autograd.scipy.signal.convolve