Gated Linear Unit (GLU) module. Args: dim_in (int): Input dimension. dim_out (int): Output dimension. activation (Callable): Activation function to be applied to the gate. mult_bias (bool, optional): Whether to multiply the bias term. Defaults to False.
| 16 | |
| 17 | # [GLU] |
| 18 | class GLU(nn.Module): |
| 19 | """ |
| 20 | Gated Linear Unit (GLU) module. |
| 21 | |
| 22 | Args: |
| 23 | dim_in (int): Input dimension. |
| 24 | dim_out (int): Output dimension. |
| 25 | activation (Callable): Activation function to be applied to the gate. |
| 26 | mult_bias (bool, optional): Whether to multiply the bias term. Defaults to False. |
| 27 | linear (Callable, optional): Linear function to be used for projection. Defaults to False. |
| 28 | """ |
| 29 | |
| 30 | def __init__( |
| 31 | self, |
| 32 | dim_in: int, |
| 33 | dim_out: int, |
| 34 | activation: Callable, |
| 35 | mult_bias: bool = False, |
| 36 | linear: Callable = False, |
| 37 | *args, |
| 38 | **kwargs |
| 39 | ): |
| 40 | super().__init__() |
| 41 | self.dim_in = dim_in |
| 42 | self.dim_out = dim_out |
| 43 | self.activation = activation |
| 44 | self.mult_bias = mult_bias |
| 45 | |
| 46 | if linear: |
| 47 | self.proj = linear(dim_in, dim_out * 2) |
| 48 | else: |
| 49 | self.proj = BitLinear(dim_in, dim_out * 4, *args, **kwargs) |
| 50 | |
| 51 | self.mult_bias = nn.Parameter(torch.ones(dim_out)) if mult_bias else 1.0 |
| 52 | |
| 53 | def forward(self, x: Tensor): |
| 54 | x, gate = self.proj(x).chunk(2, dim=-1) |
| 55 | return x * self.activation(gate) * self.mult_bias |
| 56 | |
| 57 | |
| 58 | # [FEATURE] Add type hints to the forward method |