import torch.nn as nn


class MLP(nn.Module):
    """Simple feed-forward MLP used across examples and tutorials."""

    def __init__(self, input_dim, hidden_dims, output_dim, p_drop=0.0):
        super().__init__()
        dims = [input_dim] + list(hidden_dims) + [output_dim]
        layers = []
        for i in range(len(dims) - 2):
            layers += [nn.Linear(dims[i], dims[i + 1]), nn.ReLU(), nn.Dropout(p_drop)]
        layers += [nn.Linear(dims[-2], dims[-1])]
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)
