Skip to content

Support the return_sign argument in logsumexp derivatives - #809

Open
ump45nose wants to merge 1 commit into
HIPS:masterfrom
ump45nose:fix/logsumexp-return-sign
Open

Support the return_sign argument in logsumexp derivatives#809
ump45nose wants to merge 1 commit into
HIPS:masterfrom
ump45nose:fix/logsumexp-return-sign

Conversation

@ump45nose

Copy link
Copy Markdown

scipy.special.logsumexp returns a (log(abs(sum(b * exp(x)))), sign) pair when it is
called with return_sign=True, but the derivatives of logsumexp did not accept that
argument, so both modes raised a TypeError instead of differentiating:

import autograd.scipy.special as asp
from autograd import grad

x = np.array([0.5, 1.5, 2.5])
b = np.array([1.0, -1.0, 1.0])

grad(lambda x: asp.logsumexp(x, b=b, return_sign=True)[0])(x)
# TypeError: make_grad_logsumexp() got an unexpected keyword argument 'return_sign'

This implements the request in #243 (the workaround there was to reimplement logsumexp
in Python and differentiate that).

Accepting the argument is not enough on its own: the returned value is the log of the
absolute value of the sum, so its derivative is

d/dx log(abs(sum(b * exp(x)))) = sign * b * exp(x - ans)

i.e. the sign has to appear in the gradient, otherwise the gradient points the wrong way
whenever the sum is negative. The sign itself is piecewise constant, so only the
log-sum-exp component of the pair has a non-zero derivative (and the tangent of the sign
is zero in forward mode).

Both the VJP (make_grad_logsumexp) and the JVP (fwd_grad_logsumexp) are updated; the
existing code path, where return_sign is not used, is untouched.

Tests

  • test_logsumexp_negative_sum pins the sign handling down with a case whose sum is
    exactly negative (b = [1, -3], x = [0, 0], so sign = -1 and ans = log(2)), and
    asserts the gradient equals b * exp(x) / sum(b * exp(x)) as well as checking it
    numerically.
  • test_logsumexp_return_sign runs the same shapes/axes/keepdims combinations as the
    existing logsumexp tests with return_sign=True and mixed-sign b, in both modes.

Both fail on main with the TypeError above and pass with this change. The suite is
otherwise unchanged:

$ pytest tests/ -q
634 passed, 13 skipped     # 632 passed, 13 skipped before this change

ruff and ruff-format are clean.

Closes #243

``scipy.special.logsumexp`` returns a ``(log(abs(sum(b * exp(x)))), sign)``
pair when called with ``return_sign=True``, but the derivatives of
``logsumexp`` did not accept that argument, so both modes failed with a
TypeError instead of differentiating:

    grad(lambda x: asp.logsumexp(x, b=b, return_sign=True)[0])(x)
    # TypeError: make_grad_logsumexp() got an unexpected keyword argument 'return_sign'

Accept ``return_sign`` in the VJP and the JVP. The sign is piecewise
constant, so only the log-sum-exp component of the pair has a non-zero
derivative, but it does change that derivative: since the returned value is
the log of the *absolute value* of the sum, its derivative is
``sign * b * exp(x - ans)``, and the sign has to be taken into account to
get the direction right when the sum is negative.

Closes HIPS#243
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

logsumexp for values with sign

2 participants