def count_params(input_dim, hid_dim, is_plastic) ->int:
    plastic_params = hid_dim * hid_dim + 1 if is_plastic else 0
    return 3 * (input_dim * hid_dim) + 3 * (hid_dim * hid_dim
        ) + 2 + 2 + plastic_params


print(count_params(64, 80, True))
print(count_params(64, 128, False))

1a6441c59b6b9ba26cdb8009d13fa3a72db04f2e
<mask token>

def count_params(input_dim, hid_dim, is_plastic) ->int:
    plastic_params = hid_dim * hid_dim + 1 if is_plastic else 0
    return 3 * (input_dim * hid_dim) + 3 * (hid_dim * hid_dim
        ) + 2 + 2 + plastic_params


<mask token>

def count_params(input_dim, hid_dim, is_plastic) ->int:
    plastic_params = hid_dim * hid_dim + 1 if is_plastic else 0
    return 3 * (input_dim * hid_dim) + 3 * (hid_dim * hid_dim
        ) + 2 + 2 + plastic_params


print(count_params(64, 80, True))
print(count_params(64, 128, False))
