Neural Mastery
← Practice
Vectorized Multi-Head Attention Reshape and Transposemedium

Vectorized Multi-Head Attention Reshape and Transpose

medium NumPy & Data Wrangling 15-20 min

NumPy allowed · plain Python earns 10 bonus points

Attention step 1/4

Tokens attend to tokens

Q · Kᵀ / √d → softmax → weighted values, computed in the page.
Open lesson →
01tokens
02Q K V
03scores
04attention
modelslearnrelationships
models32%
learn38%
relationships30%

Q · Kᵀ / √d scores normalized with softmax

Reshape and permute token hidden representations to isolate parallel multi-head attention streams.

Related Skills

Task

Implement split_multi_head_projections(hidden_states, num_heads, head_dim): transform a 3D tensor of shape (B, S, num_heads * head_dim) into a 4D tensor of shape (B, num_heads, S, head_dim) by splitting the last dimension into heads and swapping sequence and head dimensions.

Function Signature

split_multi_head_projections(hidden_states: list[list[list[float]]], num_heads: int, head_dim: int) -> list[list[list[list[float]]]]

Examples

Example 1: Batch 1 Seq 2 Heads 2 Dim 2
Input: hidden_states = [[[1, 2, 3, 4], [5, 6, 7, 8]]] num_heads = 2 head_dim = 2
Output: [[[[1, 2], [5, 6]], [[3, 4], [7, 8]]]]
Explanation: Split hidden dimension 4 into 2 heads of dim 2.
Example 2: Batch 2 Seq 1 Heads 2 Dim 3
Input: hidden_states = [[[1, 2, 3, 4, 5, 6]], [[7, 8, 9, 10, 11, 12]]] num_heads = 2 head_dim = 3
Output: [[[[1, 2, 3]], [[4, 5, 6]]], [[[7, 8, 9]], [[10, 11, 12]]]]
Explanation: Batch of 2 sequences split into 2 heads.

Constraints

  • B >= 1
  • S >= 1
  • num_heads >= 1
  • head_dim >= 1
Python3Saved ✓
def split_multi_head_projections(hidden_states: list[list[list[float]]], num_heads: int, head_dim: int) -> list[list[list[list[float]]]]:
"""Reshape (B, S, num_heads * head_dim) into (B, num_heads, S, head_dim).
Return 4D tensor representation."""
# Implement solution logic here
pass

Case 1: Batch 1 Seq 2 Heads 2 Dim 2
Input: hidden_states = [[[1, 2, 3, 4], [5, 6, 7, 8]]] num_heads = 2 head_dim = 2
Expected: [[[[1, 2], [5, 6]], [[3, 4], [7, 8]]]]
Case 2: Batch 2 Seq 1 Heads 2 Dim 3
Input: hidden_states = [[[1, 2, 3, 4, 5, 6]], [[7, 8, 9, 10, 11, 12]]] num_heads = 2 head_dim = 3
Expected: [[[[1, 2, 3]], [[4, 5, 6]]], [[[7, 8, 9]], [[10, 11, 12]]]]