| 173 | |
| 174 | |
| 175 | class MLP(nn.Module): |
| 176 | def __init__( |
| 177 | self, |
| 178 | input_dim: int, |
| 179 | hidden_dim: int, |
| 180 | output_dim: int, |
| 181 | num_layers: int, |
| 182 | sigmoid_output: bool = False, |
| 183 | ) -> None: |
| 184 | super().__init__() |
| 185 | self.num_layers = num_layers |
| 186 | self.proj_in = nn.Linear(input_dim, hidden_dim) |
| 187 | self.layers = [nn.Linear(hidden_dim, hidden_dim) for _ in range(num_layers)] |
| 188 | self.proj_out = nn.Linear(hidden_dim, output_dim) |
| 189 | self.sigmoid_output = sigmoid_output |
| 190 | |
| 191 | def __call__(self, x): |
| 192 | x = nn.relu(self.proj_in(x)) |
| 193 | for i, layer in enumerate(self.layers): |
| 194 | x = nn.relu(layer(x)) |
| 195 | x = self.proj_out(x) |
| 196 | if self.sigmoid_output: |
| 197 | x = mx.sigmoid(x) |
| 198 | return x |
| 199 | |
| 200 | |
| 201 | # TODO: Naive implem. Replace when mlx.nn support conv_transpose |