Cross-Attention Shapes

Problem

Compute the shapes inside encoder-decoder cross attention. With decoder length dec_len, encoder length enc_len, num_heads and head_dim, the attention weight tensor has shape [num_heads, dec_len, enc_len] and the output has shape [dec_len, num_heads, head_dim].

The signature is cross_attention_shapes(dec_len, enc_len, num_heads, head_dim). Return [weights_shape, output_shape].

Hint

Queries come from the decoder and keys/values from the encoder, so the weight matrix is dec_len x enc_len.

Test cases

case 1

in: [4,6,8,64]

expected: [[8,4,6],[4,8,64]]

case 2

in: [1,1,1,1]

expected: [[1,1,1],[1,1,1]]

case 3

in: [2,5,2,3]

expected: [[2,2,5],[2,2,3]]

case 4

in: [10,10,4,16]

expected: [[4,10,10],[10,4,16]]

Starter codePython

def cross_attention_shapes(dec_len, enc_len, num_heads, head_dim):
    # Your code here
    pass

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