Custom Autograd Function

Write your own torch.autograd.Function for a stable log-sum-exp, with a backward pass that you derive by hand.

Log-sum-exp over the last axis is:

lse(x) = log(exp(x_1) + exp(x_2) + ... + exp(x_n))

Computed directly, exp(1000) overflows to inf. Subtract the row maximum c first. The answer does not change:

lse(x) = c + log(sum(exp(x - c)))

The derivative of lse with respect to x_i is exp(x_i - lse(x)), which is the softmax of x. In backward, multiply that by the incoming gradient.

Complete the forward and backward static methods of LogSumExp. In forward, use ctx.save_for_backward to keep what backward needs.

Function signature

class LogSumExp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x): ...
    @staticmethod
    def backward(ctx, grad_output): ...

Example

Input

x = torch.tensor([[1000.0, 1000.0]], requires_grad=True)
out = logsumexp(x)
out.sum().backward()

Output

out    = tensor([1000.6931])
x.grad = tensor([[0.5000, 0.5000]])

Reasoning

Subtracting the maximum 1000 leaves log(1 + 1) = 0.6931. The softmax of two equal values is 0.5 each, and the incoming gradient of sum is 1.

Constraints and notes

• x has shape (..., n) and finite values. The output has shape (...).
• grad_output has shape (...), and backward returns a tensor of shape (..., n).
• Keep the dtype of the input. Do not call torch.logsumexp or torch.softmax.