Logsumexp Tracking Step

Problem

Merge a new value x into a running log-sum-exp lse. If x > lse, return x + log1p(exp(lse - x)), otherwise lse + log1p(exp(x - lse)). Use log1p for stability.

The signature is logsumexp_step(lse, x). Return the new log-sum-exp.

Hint

Log-sum-exp turns a product of probabilities into a stable sum in log space.

Test cases

case 1

in: [0,1]

expected: 1.3132616875

case 2

in: [0,0]

expected: 0.6931471806

case 3

in: [2,0.5]

expected: 2.201413278

case 4

in: [-1000000000,5]

expected: 5

Starter codePython

import math
def logsumexp_step(lse, x):
    # Your code here
    pass

4 test cases · run your code against them in the editor