LSTM Parameter Count for Forecasting

Problem

Compute the number of trainable parameters in a stacked LSTM. Each layer has four gates with hidden_size * (input + hidden_size + 1) weights per gate, where the input size is input_size for the first layer and hidden_size for the rest.

Return the total parameter count.

Hint

The +1 per gate is the bias term.

Test cases

case 1

in: [1,4,1]

expected: 96

case 2

in: [2,3,1]

expected: 72

case 3

in: [1,4,2]

expected: 240

case 4

in: [0,2,1]

expected: 24

Starter codePython

def lstm_param_count(input_size, hidden_size, layers):
    # Your code here
    pass

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