EasyNLPnlp-071

Attention Head Concatenation Shape

Problem

Return the shape [seq_len, n_heads * d_head] produced by concatenating n_heads attention heads of width d_head over a sequence of length seq_len.

Hint

Concatenation keeps the sequence length and sums the head widths.

case 1

in: [5,2,4]

expected: [5,8]

case 2

in: [1,1,1]

expected: [1,1]

case 3

in: [10,8,64]

expected: [10,512]

case 4

in: [0,2,3]

expected: [0,6]

Starter codePython

def head_concat_shape(seq_len, n_heads, d_head):
    # Your code here
    pass

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