| 75 | |
| 76 | class MLP(nn.Module): |
| 77 | def __init__( |
| 78 | self, |
| 79 | dim_in: int, |
| 80 | dim_out: int, |
| 81 | n_neurons: int, |
| 82 | n_hidden_layers: int, |
| 83 | activation: str = "silu", |
| 84 | output_activation: Optional[str] = "silu", |
| 85 | bias: bool = True, |
| 86 | dropout: float = 0.0, |
| 87 | use_residual: bool = False, |
| 88 | use_rmsnorm: bool = False, |
| 89 | ): |
| 90 | super().__init__() |
| 91 | self.use_residual = use_residual |
| 92 | self.use_rmsnorm = use_rmsnorm |
| 93 | self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() |
| 94 | input_norm = LayerNorm(dim_in) if use_rmsnorm else nn.Identity() |
| 95 | output_norm = nn.Identity() # no normalization for the output |
| 96 | |
| 97 | layers = [ |
| 98 | input_norm, |
| 99 | self.make_linear( |
| 100 | dim_in, n_neurons, is_first=True, is_last=False, bias=bias |
| 101 | ), |
| 102 | self.make_activation(activation), |
| 103 | self.dropout, |
| 104 | ] |
| 105 | for i in range(n_hidden_layers - 1): |
| 106 | layers += [ |
| 107 | self.make_linear( |
| 108 | n_neurons, n_neurons, is_first=False, is_last=False, bias=bias |
| 109 | ), |
| 110 | self.make_activation(activation), |
| 111 | self.dropout, |
| 112 | ] |
| 113 | layers += [ |
| 114 | self.make_linear( |
| 115 | n_neurons, dim_out, is_first=False, is_last=True, bias=bias |
| 116 | ), |
| 117 | output_norm, |
| 118 | ] |
| 119 | self.layers = nn.Sequential(*layers) |
| 120 | self.output_activation = self.make_activation(output_activation) |
| 121 | |
| 122 | def forward(self, x): |
| 123 | if self.use_residual: |