SwiGLU FFN Parameter Count

Problem

Count the parameters of a SwiGLU feed-forward block without biases: three matrices of shape (d_ff, d_model) give 3 * d_model * d_ff. When bias is True, add one bias vector of length d_ff for each of the two up projections and one of length d_model for the down projection.

The signature is swiglu_ffn_params(d_model, d_ff, bias=False). Return an integer.

Hint

SwiGLU uses three weight matrices because the gate and value branches are separate.

Test cases

case 1

in: [512,1365]

expected: 2096640

case 2

in: [768,2048]

expected: 4718592

case 3

in: [512,1365,true]

expected: 2099882

case 4

in: [64,128]

expected: 24576

Starter codePython

def swiglu_ffn_params(d_model, d_ff, bias=False):
    # Your code here
    pass

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