Streaming Entropy, Part 2: Numerically Stable Entropy

Streaming Entropy Interview Sequence
This is Part 2 of one four-part ML coding exercise.
Streaming Entropy, Part 1: Batch Entropy
Streaming Entropy, Part 2: Numerically Stable Entropy
Streaming Entropy, Part 3: Block-wise Entropy
Streaming Entropy, Part 4: Stable Streaming Entropy
As in Part 1, logits define a softmax distribution:
p[i] = exp(x[i]) / sum(exp(x[j]))
and its entropy is
H = -sum(p[i] * ln(p[i]))
This time the logits may be extremely large or small. Directly computing exp(x[i]) can overflow, and computing log(exp(x[i]) / sumExp) can take the logarithm of a probability that rounded to zero.
Numerical Stability Requirements
Return the non-negative entropy without overflow, underflow-driven NaN, or loss of the log-softmax term.
Use a maximum shift for the exponentials and derive log-softmax algebraically from the shifted logits and their normalization sum. Your implementation should run in O(n) time.
An answer is accepted when its absolute or relative error is at most 1e-6.
Continue to Part 3: Block-wise Entropy.

Function
stableEntropy(logits: double[]) → double

Examples
Example 1
logits = [1000.0, 1000.0, -1000.0]
return = 0.6931471805599453
Directly evaluating exp(1000) overflows. After shifting by the maximum, the third probability is negligible and the first two are each effectively 0.5.
Example 2
logits = [999999998.0, 999999999.0, 1000000000.0]
return = 0.8323955818399389
Subtracting the maximum turns the logits into [-2, -1, 0] without changing the softmax distribution.
Example 3
logits = [42.0]
return = 0.0
A one-element softmax distribution is [1], whose entropy is zero.

Constraints
1 <= logits.length <= 2 * 10^5
Every logit is finite and lies in [-10^9, 10^9].
Use the natural logarithm.
Use O(n) time.
The result must remain finite for every valid input.

The trick is log-sum-exp. Let m = max(x), s = sum(exp(x[i] - m)), and L = ln(s). Then log p[i] = (x[i] - m) - L and p[i] = exp(x[i] - m) / s. Entropy is H = -sum(p[i] * logp[i]). Every exponent is at most 0, so nothing overflows, and the max element contributes exp(0) = 1, so s is at least 1 and L is never undefined. The common pitfall is computing log(p[i]) after p[i] has underflowed to zero, which gives -inf times 0 and a NaN. Don't take the log of p. Use the algebraic form instead. Example 1 shows it: the -1000 logit gives p near 0, but its log-softmax term is finite, so the product vanishes cleanly. Single-element input returns 0. Watch for -0.0 and clamp tiny negatives to 0. StealthCoder is your hedge if the live OA rattles you, but this is a three-line formula once you see it.

The honest play: practice the pattern, and have StealthCoder ready for the one you didn't see coming.