Log-softmax

Background

Log-softmax computes $\log(\text{softmax}(x))$ in a numerically stable way. It appears everywhere softmax and log show up together — cross-entropy loss, log-likelihoods, beam search — because computing them jointly avoids the overflow of exp on large logits and the $\log 0$ underflow on tiny probabilities.

Problem statement

Implement log_softmax(x) over the last axis:

$$\text{log\_softmax}(x)_i = x_i - \max_j x_j - \log\sum_j e^{\,x_j - \max_k x_k}$$

The max-subtraction is the stability trick; it cancels mathematically but prevents overflow.

Input

• x — np.ndarray: 1-D logits, or 2-D (batch, classes) (softmax taken over the last axis).

Output

Returns an np.ndarray of the same shape: the log-probabilities (each row's exp sums to 1).

Examples

Example 1

Input:  x = [1, 2, 3]
Output: [-2.4076, -1.4076, -0.4076]

Explanation: subtract the max ($3$) to get $[-2, -1, 0]$, then subtract $\log(e^{-2}+e^{-1}+e^{0}) \approx 0.4076$.

Constraints

• Subtract the per-row max before exponentiating (numerical stability).
• Operate over the last axis; support both 1-D and 2-D inputs.
• exp(log_softmax(x)) must sum to 1 along the last axis; tests use atol=1e-6.

Notes

• $\text{log\_softmax} = x - \text{logsumexp}(x)$; never compute log(softmax(x)) as two separate steps — the intermediate softmax can underflow to 0.
• Cross-entropy loss is just $-\text{log\_softmax}(\text{logits})[\text{correct class}]$, which is why frameworks fuse the two operations.