diff --git a/autograd/scipy/special.py b/autograd/scipy/special.py index 2e435900..423ce487 100644 --- a/autograd/scipy/special.py +++ b/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) diff --git a/tests/test_scipy.py b/tests/test_scipy.py index 503b618d..1f2e6d8c 100644 --- a/tests/test_scipy.py +++ b/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