Interleaved Head Split

Problem

Split a (seq_len, d_model) matrix into num_heads heads using an interleaved column layout: head h takes columns h, h + num_heads, h + 2 * num_heads, and so on.

The signature is split_heads_interleaved(x, num_heads). This differs from the contiguous chunk split used by most frameworks.

Hint

Head h owns the columns whose index is congruent to h modulo num_heads.

Test cases

case 1

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

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

case 2

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

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

case 3

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

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

case 4

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

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

Starter codePython

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

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