| """Small neural-network layers used by RiboSphere.""" | |
| from __future__ import annotations | |
| from torch import Tensor, nn | |
| class FeedForward(nn.Module): | |
| """Two-layer MLP equivalent to the subset of timm.layers.Mlp used here.""" | |
| def __init__( | |
| self, | |
| input_features: int, | |
| hidden_features: int, | |
| output_features: int, | |
| *, | |
| activation: type[nn.Module] = nn.GELU, | |
| dropout: float = 0.0, | |
| ) -> None: | |
| super().__init__() | |
| if min(input_features, hidden_features, output_features) <= 0: | |
| raise ValueError("All feature dimensions must be positive.") | |
| if not 0.0 <= dropout < 1.0: | |
| raise ValueError("dropout must be in [0, 1).") | |
| self.fc1 = nn.Linear(input_features, hidden_features) | |
| self.activation = activation() | |
| self.dropout1 = nn.Dropout(dropout) | |
| self.fc2 = nn.Linear(hidden_features, output_features) | |
| self.dropout2 = nn.Dropout(dropout) | |
| def forward(self, inputs: Tensor) -> Tensor: | |
| """Apply the MLP without changing leading dimensions.""" | |
| outputs = self.fc1(inputs) | |
| outputs = self.activation(outputs) | |
| outputs = self.dropout1(outputs) | |
| outputs = self.fc2(outputs) | |
| outputs = self.dropout2(outputs) | |
| return outputs | |