logsumexp: Numerically Stable Log-Sum-Exp in Python

import numpy as np

# Large values: exp() overflows to inf
x = np.array([1000, 1001, 1002])
print(np.exp(x))          # => [inf inf inf]
print(np.log(np.sum(np.exp(x))))  # => nan

# Very negative values: exp() underflows to 0
x = np.array([-1000, -1001, -1002])
print(np.exp(x))          # => [0. 0. 0.]
print(np.log(np.sum(np.exp(x))))  # => -inf (wrong!)

def logsumexp_manual(x):
    c = np.max(x)
    return c + np.log(np.sum(np.exp(x - c)))

x = np.array([1000, 1001, 1002])
print(logsumexp_manual(x))  # => 1002.4076059644443  (correct!)

from scipy.special import logsumexp
import numpy as np

# 1D array
a = np.array([1, 2, 3])
result = logsumexp(a)
print(result)  # => 3.4076059644443804
# Verification: log(exp(1) + exp(2) + exp(3)) = 3.4076...

# 2D array — along axis=1 (row-wise)
a = np.array([[1, 2, 3],
              [4, 5, 6]])

print(logsumexp(a, axis=1))
# => [3.40760596 6.40760596]
# Row 0: log(exp(1) + exp(2) + exp(3)) = 3.4076...
# Row 1: log(exp(4) + exp(5) + exp(6)) = 6.4076...

# Along axis=0 (column-wise)
print(logsumexp(a, axis=0))
# => [4.31326169 5.31326169 6.31326169]

from scipy.special import logsumexp

a = np.array([1.0, 2.0, 3.0])
b = np.array([0.5, 1.0, 2.0])  # weights

result = logsumexp(a, b=b)
# = log(0.5*exp(1) + 1.0*exp(2) + 2.0*exp(3))
print(result)  # => 3.7194...

result, sign = logsumexp(a, b=[-1, 1, -1], return_sign=True)
# result = log(|sum(b * exp(a))|)
# sign   = sign of the sum (+1 or -1)
print(result, sign)

def log_softmax(x):
    """Numerically stable log-softmax."""
    return x - logsumexp(x)

def softmax(x):
    """Numerically stable softmax."""
    return np.exp(log_softmax(x))

logits = np.array([2.0, 1.0, 0.1])
print(softmax(logits))
# => [0.65900114 0.24243297 0.09856589]
print(softmax(logits).sum())  # => 1.0

# Log-probabilities of each observation
log_probs = np.array([-0.5, -1.2, -0.8, -2.1, -0.3])

# Total log-likelihood (sum in log space)
total_log_likelihood = np.sum(log_probs)
print(total_log_likelihood)  # => -4.9

# When combining probabilities from different models (mixture):
# log(p1 * w1 + p2 * w2) = logsumexp([log_p1 + log_w1, log_p2 + log_w2])
log_p1, log_p2 = -0.5, -1.0
log_w1, log_w2 = np.log(0.7), np.log(0.3)

log_mixture = logsumexp([log_p1 + log_w1, log_p2 + log_w2])
print(log_mixture)

def forward_step(log_alpha, log_transition, log_emission):
    """One step of the forward algorithm in log space."""
    # log_alpha: shape (n_states,)
    # log_transition: shape (n_states, n_states)
    # log_emission: shape (n_states,)

    # For each next state j: log(sum_i alpha_i * T_ij) + log(E_j)
    log_alpha_next = logsumexp(
        log_alpha[:, np.newaxis] + log_transition,
        axis=0
    ) + log_emission

    return log_alpha_next

# Unnormalized log-probabilities
log_unnorm = np.array([-1.0, -2.0, -0.5, -3.0])

# Normalize: subtract log(sum(exp(log_unnorm)))
log_norm = log_unnorm - logsumexp(log_unnorm)

# Verify they sum to 1 in probability space
print(np.exp(log_norm).sum())  # => 1.0
print(log_norm)
# => [-1.56...  -2.56...  -1.06...  -3.56...]

def logsumexp_numpy(a, axis=None):
    """NumPy-only logsumexp."""
    a_max = np.max(a, axis=axis, keepdims=True)
    out = np.log(np.sum(np.exp(a - a_max), axis=axis))
    out += np.squeeze(a_max, axis=axis)
    return out

# Test
a = np.array([[1, 2, 3], [4, 5, 6]])
print(logsumexp_numpy(a, axis=1))
# => [3.40760596 6.40760596]