| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| def _round_to_multiple(x: int, multiple_of: int) -> int: |
| multiple_of = int(multiple_of) |
| if multiple_of <= 0: |
| return int(x) |
| x = int(x) |
| return int(((x + multiple_of - 1) // multiple_of) * multiple_of) |
|
|
|
|
| class VisualSwiGLUFFN(nn.Module): |
| """ |
| Plain pointwise SwiGLU FFN: |
| |
| y = W_o( silu(g) * v ), where [g, v] = W_proj(x) |
| |
| No convolutions, no spatial/temporal mixing, no grid reshapes. Operates on the |
| last dim only, so it accepts any token layout ([B, L, C], [B, F, HW, C], ...) |
| and returns the same shape. Mixing between tokens is left entirely to attention. |
| """ |
|
|
| def __init__( |
| self, |
| dim: int, |
| mlp_ratio: float = 4.0, |
| multiple_of: int = 16, |
| bias: bool = True, |
| ): |
| super().__init__() |
| base_hidden = int(dim * mlp_ratio) |
| hidden = int((2.0 / 3.0) * base_hidden) |
| hidden = _round_to_multiple(hidden, multiple_of) |
|
|
| self.hidden_dim = hidden |
| self.proj = nn.Linear(dim, 2 * hidden, bias=bias) |
| self.out = nn.Linear(hidden, dim, bias=bias) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| g, v = self.proj(x).chunk(2, dim=-1) |
| return self.out(F.silu(g) * v) |
|
|
