Implement a **numerically stable softmax** as used inside Transformer attention, so you can safely turn attention logits into probabilities without overflow. The softmax is defined as:
$$
\text{softmax}(z_i)=\frac{e^{z_i}}{\sum_{j=1}^{n} e^{z_j}}
$$

## Requirements

Implement the function
```python
import numpy as np

def stable_softmax(logits):
    \"""
    Computes a numerically stable softmax over the last dimension.

    Args:
        logits (np.ndarray): 2D attention logits of shape (num_queries, num_keys).

    Returns:
        np.ndarray: Softmax probabilities with the same shape as logits.
    \"""
    # Your solution below
```

**Rules:**
- Compute softmax **row-wise**: each query’s logits should sum to 1 across keys.
- Use the stability trick: subtract the row max before exponentiating.
- Use only NumPy operations (no SciPy / deep learning frameworks).
- Return the output as a NumPy array.
- Keep it numerically stable for large positive/negative logits.

## Example

```python
logits = [
    [1000.0, 1000.0, 999.0],
    [-1000.0, 0.0, 1000.0]
]

stable_softmax(logits)
```

Output:
```python
[
    [0.4223, 0.4223, 0.1554],
    [0.0, 0.0, 1.0]
]
```

Baseball Game

Top City by Rides

Stable softmax

Solve this easy level ML coding question from General. The premier LeetCode for ML platform to master Python, SQL, and Machine Learning algorithms.

1131. Stable softmax

Examples

Constraints

Similar Questions