Implement a **single-head Transformer self-attention block** (forward pass only) that takes a sequence of token embeddings and returns the attention-mixed outputs. This question tests whether you understand scaled dot-product attention and can implement it cleanly using basic NumPy ops.

The attention is defined as:
$$
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d}}\right)V
$$

## Requirements

Implement the function
```python
import numpy as np

def self_attention_block(X):
    \"""
    Computes the forward pass of a single-head self-attention block.

    Args:
        X (np.ndarray): Input token embeddings with shape (seq_len, d).

    Returns:
        np.ndarray: Attention output with shape (seq_len, d).
    \"""
    # Your solution below
```

**Rules:**
- Treat `Q = K = V = X` (no learned projection matrices in this question).
- Compute scaled dot-product attention scores: `S = (Q @ K.T) / sqrt(d)`.
- Apply a **numerically stable softmax** across each row of `S` to get attention weights.
- Compute the output as `A @ V`, where `A` is the attention weight matrix.
- Return the result as a **NumPy array**.

## Example

```python
X = np.array([
    [1.0, 0.0],
    [0.0, 1.0]
])

self_attention_block(X)
```

Output:
```python
array([
    [0.7311, 0.2689],
    [0.2689, 0.7311]
])
```

Streaming Median Finder

Stream Deduplication Engine

Transformer self attention block

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

159. Transformer self attention block

Examples

Constraints

Similar Questions