(
self,
dim: int,
dim_out: int | None = None,
mult: int = 4,
dropout: float = 0.0,
activation_fn: str = "geglu",
final_dropout: bool = False,
inner_dim=None,
bias: bool = True,
)
| 1694 | """ |
| 1695 | |
| 1696 | def __init__( |
| 1697 | self, |
| 1698 | dim: int, |
| 1699 | dim_out: int | None = None, |
| 1700 | mult: int = 4, |
| 1701 | dropout: float = 0.0, |
| 1702 | activation_fn: str = "geglu", |
| 1703 | final_dropout: bool = False, |
| 1704 | inner_dim=None, |
| 1705 | bias: bool = True, |
| 1706 | ): |
| 1707 | super().__init__() |
| 1708 | if inner_dim is None: |
| 1709 | inner_dim = int(dim * mult) |
| 1710 | dim_out = dim_out if dim_out is not None else dim |
| 1711 | |
| 1712 | if activation_fn == "gelu": |
| 1713 | act_fn = GELU(dim, inner_dim, bias=bias) |
| 1714 | if activation_fn == "gelu-approximate": |
| 1715 | act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias) |
| 1716 | elif activation_fn == "geglu": |
| 1717 | act_fn = GEGLU(dim, inner_dim, bias=bias) |
| 1718 | elif activation_fn == "geglu-approximate": |
| 1719 | act_fn = ApproximateGELU(dim, inner_dim, bias=bias) |
| 1720 | elif activation_fn == "swiglu": |
| 1721 | act_fn = SwiGLU(dim, inner_dim, bias=bias) |
| 1722 | elif activation_fn == "linear-silu": |
| 1723 | act_fn = LinearActivation(dim, inner_dim, bias=bias, activation="silu") |
| 1724 | |
| 1725 | self.net = nn.ModuleList([]) |
| 1726 | # project in |
| 1727 | self.net.append(act_fn) |
| 1728 | # project dropout |
| 1729 | self.net.append(nn.Dropout(dropout)) |
| 1730 | # project out |
| 1731 | self.net.append(nn.Linear(inner_dim, dim_out, bias=bias)) |
| 1732 | # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout |
| 1733 | if final_dropout: |
| 1734 | self.net.append(nn.Dropout(dropout)) |
| 1735 | |
| 1736 | def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: |
| 1737 | if len(args) > 0 or kwargs.get("scale", None) is not None: |
no test coverage detected