MLP FLOPs

Problem

Count the multiply-add FLOPs of a forward pass through a 2-layer MLP, treating one multiply-add as 2 FLOPs: 2 * input_dim * hidden_dim + 2 * hidden_dim * output_dim.

The signature is mlp_flops(input_dim, hidden_dim, output_dim). Return an integer.

Hint

Each layer contributes twice the product of its input and output sizes; biases are ignored.

Test cases

case 1

in: [4,8,2]

expected: 96

case 2

in: [784,128,10]

expected: 203264

case 3

in: [1,1,1]

expected: 4

case 4

in: [3,3,3]

expected: 36

Starter codePython

def mlp_flops(input_dim, hidden_dim, output_dim):
    # Your code here
    pass

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