Multi-Head Split

Problem

Split a (seq_len, d_model) matrix into num_heads attention heads of shape (seq_len, d_model / num_heads) by slicing contiguous column chunks.

The signature is split_heads(x, num_heads). Return a list of head matrices; head h uses columns [h * head_dim, (h + 1) * head_dim).

Hint

Head h takes the contiguous column block starting at h * (d_model // num_heads).

case 1

in: [[[1,2,3,4],[5,6,7,8]],2]

expected: [[[1,2],[5,6]],[[3,4],[7,8]]]

case 2

in: [[[1,2],[3,4]],1]

expected: [[[1,2],[3,4]]]

case 3

in: [[[1,2,3,4]],4]

expected: [[[1]],[[2]],[[3]],[[4]]]

case 4

in: [[[1,2,3,4,5,6]],3]

expected: [[[1,2]],[[3,4]],[[5,6]]]

Starter codePython

def split_heads(x, num_heads):
    # Your code here
    pass

4 test cases · run your code against them in the editor