Implement Stable Softmax

Problem Statement

The softmax function is a critical component in machine learning, mapping a vector of real numbers to a probability distribution. Given a vector logits of size N, the softmax function is defined as P(i) = exp(logits[i]) / sum(exp(logits[j])) for all j in 0 to N-1.

A naive implementation of softmax is prone to numerical instability. Specifically, exp(x) grows extremely rapidly, and for large values of logits[i], the calculation can easily result in an arithmetic overflow (inf), leading to NaN outputs.

Implement the stable_softmax function. You must ensure the result remains accurate even when input values are very large.

Function Signature

def stable_softmax(logits: list[float]) -> list[float]:
    # Implementation here

Constraints

• 1 <= len(logits) <= 10^5
• -10^4 <= logits[i] <= 10^4

Example 1

Input: [1.0, 2.0, 3.0] Output: [0.090, 0.245, 0.665]

Example 2

Input: [1000.0, 1000.0] Output: [0.5, 0.5]

Note: A naive implementation using exp(1000.0) will return NaN due to overflow, but the stable version should correctly return [0.5, 0.5].