Numerically Stable Softmax

Softmax turns a vector of real-valued scores (logits) into a probability distribution. It sits at the end of every language model, turning scores over the vocabulary into next-token probabilities, and inside every attention layer.

\mathrm{softmax}(x)_i = \frac{e^{x_i}}{\sum_j e^{x_j}}

Implement softmax(x, axis=-1) with NumPy.

Requirements

Works on arrays of any shape, normalizing along axis. The output has the same shape as x.
Numerically stable: softmax(np.array([1000.0, 1001.0])) must not return nan. The naive formula overflows, because $e^{1000}$ is inf in float64.
No Python loops over elements.

Hint

Softmax is invariant to shifting all the inputs by the same constant $c$ :

\frac{e^{x_i - c}}{\sum_j e^{x_j - c}} = \frac{e^{-c}\, e^{x_i}}{e^{-c} \sum_j e^{x_j}}

Pick $c$ so the largest exponent is $e^0 = 1$ . Then nothing can overflow. Look up the keepdims argument to make the shift broadcast correctly.