GRU Parameter Count

Problem

Count the parameters of a GRU layer: 3 * (input_dim * hidden_dim + hidden_dim * hidden_dim), plus 3 * hidden_dim if bias is enabled (one bias per gate).

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

Hint

The three gates are reset, update and candidate; GRUs have one fewer gate than LSTMs.

Test cases

case 1

in: [512,512]

expected: 1574400

case 2

in: [100,128,false]

expected: 87552

case 3

in: [1,1]

expected: 9

case 4

in: [64,32]

expected: 9312

Starter codePython

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

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