Numerically Stable Softmax

Implement softmax in a way that does not overflow on large inputs.

Signature: def softmax(x: np.ndarray) -> np.ndarray

Example:

Requirements:

Output is a valid probability distribution (non-negative, sums to 1)
Must produce correct, finite results for inputs containing very large values (e.g., [1000.0, 1001.0, 1002.0]) — a naive exp(x) / exp(x).sum() returns nan here

Hint: softmax(x) is shift-invariant: softmax(x) = softmax(x - c) for any constant c. Choose c so exp cannot overflow.

Math

softmax (x_{i}) = \frac{e ^{x_{i} - m a x (x)}}{\sum _{j} e ^{x_{j} - m a x (x)}}

1. Numerically Stable Softmax