(self, dim, hidden_dim)
| 191 | |
| 192 | class MLP(nn.Module): |
| 193 | def __init__(self, dim, hidden_dim): |
| 194 | super().__init__() |
| 195 | self.gate_proj = nn.Linear(dim, hidden_dim, bias=False) |
| 196 | self.down_proj = nn.Linear(hidden_dim, dim, bias=False) |
| 197 | self.up_proj = nn.Linear(dim, hidden_dim, bias=False) |
| 198 | |
| 199 | def __call__(self, x) -> mx.array: |
| 200 | return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x)) |