| 1156 | """ |
| 1157 | |
| 1158 | def __init__( |
| 1159 | self, |
| 1160 | dim: int, |
| 1161 | dim_out: Optional[int] = None, |
| 1162 | mult: int = 4, |
| 1163 | dropout: float = 0.0, |
| 1164 | activation_fn: str = "geglu", |
| 1165 | final_dropout: bool = False, |
| 1166 | inner_dim=None, |
| 1167 | bias: bool = True, |
| 1168 | ): |
| 1169 | super().__init__() |
| 1170 | if inner_dim is None: |
| 1171 | inner_dim = int(dim * mult) |
| 1172 | dim_out = dim_out if dim_out is not None else dim |
| 1173 | |
| 1174 | if activation_fn == "gelu": |
| 1175 | act_fn = GELU(dim, inner_dim, bias=bias) |
| 1176 | if activation_fn == "gelu-approximate": |
| 1177 | act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias) |
| 1178 | elif activation_fn == "geglu": |
| 1179 | act_fn = GEGLU(dim, inner_dim, bias=bias) |
| 1180 | elif activation_fn == "geglu-approximate": |
| 1181 | act_fn = ApproximateGELU(dim, inner_dim, bias=bias) |
| 1182 | elif activation_fn == "swiglu": |
| 1183 | act_fn = SwiGLU(dim, inner_dim, bias=bias) |
| 1184 | |
| 1185 | self.net = nn.ModuleList([]) |
| 1186 | # project in |
| 1187 | self.net.append(act_fn) |
| 1188 | # project dropout |
| 1189 | self.net.append(nn.Dropout(dropout)) |
| 1190 | # project out |
| 1191 | self.net.append(nn.Linear(inner_dim, dim_out, bias=bias)) |
| 1192 | # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout |
| 1193 | if final_dropout: |
| 1194 | self.net.append(nn.Dropout(dropout)) |
| 1195 | |
| 1196 | def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: |
| 1197 | if len(args) > 0 or kwargs.get("scale", None) is not None: |