LSTM Parameter Count

Problem

Count the parameters of an LSTM layer: 4 * (input_dim * hidden_dim + hidden_dim * hidden_dim), plus 4 * hidden_dim if bias is enabled (one bias per gate).

The signature is lstm_param_count(input_dim, hidden_dim, bias=True). Return an integer.

Hint

The four gates are forget, input, cell candidate and output.

Test cases

case 1

in: [512,512]

expected: 2099200

case 2

in: [100,128,false]

expected: 116736

case 3

in: [1,1]

expected: 12

case 4

in: [64,32]

expected: 12416

Starter codePython

def lstm_param_count(input_dim, hidden_dim, bias=True):
    # Your code here
    pass

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