BitFeedForward module performs feed-forward operations on the input tensor. Args: dim (int): The input dimension. dim_out (int, optional): The output dimension. If not provided, it is set to the input dimension. mult (int, optional): The multiplier for the inner dim
| 57 | |
| 58 | # [FEATURE] Add type hints to the forward method |
| 59 | class BitFeedForward(nn.Module): |
| 60 | """ |
| 61 | BitFeedForward module performs feed-forward operations on the input tensor. |
| 62 | |
| 63 | Args: |
| 64 | dim (int): The input dimension. |
| 65 | dim_out (int, optional): The output dimension. If not provided, it is set to the input dimension. |
| 66 | mult (int, optional): The multiplier for the inner dimension. Default is 4. |
| 67 | glu (bool, optional): Whether to use Gated Linear Unit (GLU) activation. Default is False. |
| 68 | glu_mult_bias (bool, optional): Whether to apply bias to the GLU activation. Default is False. |
| 69 | swish (bool, optional): Whether to use Swish activation. Default is False. |
| 70 | relu_squared (bool, optional): Whether to use squared ReLU activation. Default is False. |
| 71 | post_act_ln (bool, optional): Whether to apply Layer Normalization after activation. Default is False. |
| 72 | dropout (float, optional): The dropout probability. Default is 0.0. |
| 73 | no_bias (bool, optional): Whether to exclude bias in linear layers. Default is False. |
| 74 | zero_init_output (bool, optional): Whether to initialize the last linear layer to 0. Default is False. |
| 75 | """ |
| 76 | |
| 77 | def __init__( |
| 78 | self, |
| 79 | dim: int, |
| 80 | dim_out: Optional[int] = None, |
| 81 | mult: int = 4, |
| 82 | glu: bool = False, |
| 83 | glu_mult_bias: bool = False, |
| 84 | swish: bool = False, |
| 85 | post_act_ln: bool = False, |
| 86 | dropout: float = 0.0, |
| 87 | no_bias: bool = False, |
| 88 | zero_init_output: bool = False, |
| 89 | *args, |
| 90 | **kwargs |
| 91 | ): |
| 92 | super().__init__() |
| 93 | inner_dim = int(dim * mult) |
| 94 | dim_out = default(dim_out, dim) |
| 95 | |
| 96 | if swish: |
| 97 | activation = nn.SiLU() |
| 98 | else: |
| 99 | activation = nn.GELU() |
| 100 | |
| 101 | if glu: |
| 102 | project_in = GLU(dim, inner_dim, activation, mult_bias=glu_mult_bias) |
| 103 | else: |
| 104 | project_in = nn.Sequential( |
| 105 | BitLinear(dim, inner_dim, bias=not no_bias, *args, **kwargs), activation |
| 106 | ) |
| 107 | if post_act_ln: |
| 108 | self.ff = nn.Sequential( |
| 109 | project_in, |
| 110 | nn.LayerNorm(inner_dim), |
| 111 | nn.Dropout(dropout), |
| 112 | BitLinear(inner_dim, dim_out, bias=not no_bias, *args, **kwargs), |
| 113 | ) |
| 114 | else: |
| 115 | self.ff = nn.Sequential( |
| 116 | project_in, |
no outgoing calls