Implementation of Log Softmax Function

Problem

Implement the log softmax function. Given a 1D array of logits, compute the log of the softmax probabilities in a numerically stable way.

Solution

import numpy as np

def log_softmax(logits):
    logits = np.array(logits, dtype=np.float64)
    max_val = np.max(logits)
    shifted = logits - max_val
    log_sum_exp = max_val + np.log(np.sum(np.exp(shifted)))
    return (logits - log_sum_exp).tolist()

Explanation

1. Subtract the maximum value from all logits for numerical stability (prevents overflow in exp).
2. Compute log(sum(exp(shifted))) and add back the max to get log_sum_exp.
3. The log softmax of each element is logit - log_sum_exp.
4. This avoids computing the softmax probabilities directly, which can underflow.

Complexity

• Time: O(n) where n is the number of logits
• Space: O(n) for the output array