Live Coding: Numerically Safe Softmax

Numerically Stable Softmax & Cross-Entropy is a classic coding problem testing an engineer's mastery over IEEE 754 floating-point limits and numerical stability tricks; naive $\text{softmax}(z)_i = \frac{e^{z_i}}{\sum e^{z_j}}$ triggers floating-point overflow (`inf`) when logits are large ($z_i = 1000$), yielding `NaN` quotients; the numerically safe algorithm exploits translation invariance: $\text{softmax}(z) = \text{softmax}(z - c)$; setting $c = \max(z)$ shifts all exponents into $(-\infty, 0]$ where the maximum value is $\exp(0)=1$, guaranteeing no overflow; furthermore, combining it with Log-Sum-Exp prevents underflow when evaluating log-probabilities in Cross-Entropy.

Low-level neural operator implementations, classification output layers, and debugging mysterious NaN loss crashes during training.

Naive implementations blow up into NaNs upon large logit inputs; shift-invariance guarantees exact mathematical equivalence and rock-solid numerical stability across extreme floating-point dynamic ranges.

Write the complete Pure Numpy implementation of numerically stable Softmax supporting arbitrary batch axes?

Mathematically prove the translation invariance of the Softmax operator for any scalar constant $c$?

Why does evaluating $\log(\text{softmax}(z))$ separately cause underflow, and how does the Log-Sum-Exp (LSE) trick resolve it?

Derive the asymptotic output distribution when temperature $\tau \to 0$ (argmax one-hot) and $\tau \to \infty$ (uniform distribution)?

Why does PyTorch's `CrossEntropyLoss` fuse LogSoftmax and NLLLoss inside a single fused CUDA kernel for numerical stability and speed?