r""" A feed-forward layer. Parameters: dim (`int`): The number of channels in the input. dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. mult (`int`, *optional*, defaults to 4): The multiplier to use for the h
| 181 | |
| 182 | |
| 183 | class FeedForward(nn.Module): |
| 184 | r""" |
| 185 | A feed-forward layer. |
| 186 | |
| 187 | Parameters: |
| 188 | dim (`int`): The number of channels in the input. |
| 189 | dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. |
| 190 | mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. |
| 191 | dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. |
| 192 | activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. |
| 193 | final_dropout (`bool` *optional*, defaults to False): Apply a final dropout. |
| 194 | """ |
| 195 | |
| 196 | def __init__( |
| 197 | self, |
| 198 | dim: int, |
| 199 | dim_out: Optional[int] = None, |
| 200 | mult: int = 4, |
| 201 | dropout: float = 0.0, |
| 202 | activation_fn: str = "geglu", |
| 203 | final_dropout: bool = False, |
| 204 | ): |
| 205 | super().__init__() |
| 206 | inner_dim = int(dim * mult) |
| 207 | dim_out = dim_out if dim_out is not None else dim |
| 208 | |
| 209 | if activation_fn == "gelu": |
| 210 | act_fn = GELU(dim, inner_dim) |
| 211 | if activation_fn == "gelu-approximate": |
| 212 | act_fn = GELU(dim, inner_dim, approximate="tanh") |
| 213 | elif activation_fn == "geglu": |
| 214 | act_fn = GEGLU(dim, inner_dim) |
| 215 | elif activation_fn == "geglu-approximate": |
| 216 | act_fn = ApproximateGELU(dim, inner_dim) |
| 217 | |
| 218 | self.net = nn.ModuleList([]) |
| 219 | # project in |
| 220 | self.net.append(act_fn) |
| 221 | # project dropout |
| 222 | self.net.append(nn.Dropout(dropout)) |
| 223 | # project out |
| 224 | self.net.append(nn.Linear(inner_dim, dim_out)) |
| 225 | # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout |
| 226 | if final_dropout: |
| 227 | self.net.append(nn.Dropout(dropout)) |
| 228 | |
| 229 | def forward(self, hidden_states): |
| 230 | for module in self.net: |
| 231 | hidden_states = module(hidden_states) |
| 232 | return hidden_states |
| 233 | |
| 234 | |
| 235 | class GELU(nn.Module): |