Multi-Head Merge

Problem

Merge a list of attention heads, each shaped (seq_len, head_dim), back into one (seq_len, d_model) matrix by concatenating their columns in order.

The signature is merge_heads(heads). This is the inverse of split_heads.

Hint

For each position concatenate the same row from every head in head order.

case 1

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

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

case 2

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

expected: [[1,2,3]]

case 3

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

expected: [[1,2,3]]

case 4

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

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

Starter codePython

def merge_heads(heads):
    # Your code here
    pass

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