Attention Output Projection FLOPs

Problem

Count the FLOPs of the attention output projection, treating a multiply-add as 2 FLOPs: 2 * batch_size * seq_len * d_model * d_model.

The signature is attention_out_proj_flops(batch_size, seq_len, d_model). Return an integer.

Hint

One output projection is one dense matrix multiply per token: 2 * d_model^2 FLOPs each.

case 1

in: [1,128,512]

expected: 67108864

case 2

in: [8,512,768]

expected: 4831838208

case 3

in: [1,1,64]

expected: 8192

case 4

in: [2,32,128]

expected: 2097152

Starter codePython

def attention_out_proj_flops(batch_size, seq_len, d_model):
    # Your code here
    pass

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