"""A module that defines utilities for building multi-layer perceptrons.
"""
from __future__ import annotations
import typing
import torch.nn as nn


[docs]def make_mlp(
    input_size: int,
    sizes: typing.List[int],
    hidden_activation: str | None = "ReLU",
    output_activation: str | None = "ReLU",
    layer_norm: bool = False,
) -> nn.Sequential:
    """Construct an MLP with specified fully-connected layers."""
    layers = []
    n_layers = len(sizes)
    sizes = [input_size] + sizes
    # Hidden layers
    for i in range(n_layers - 1):
        layers.append(nn.Linear(sizes[i], sizes[i + 1]))
        if layer_norm:
            layers.append(nn.LayerNorm(sizes[i + 1]))
        if hidden_activation is not None:
            layers.append(getattr(nn, hidden_activation)())
    # Final layer
    layers.append(nn.Linear(sizes[-2], sizes[-1]))
    if output_activation is not None:
        if layer_norm:
            layers.append(nn.LayerNorm(sizes[-1]))
        layers.append(getattr(nn, output_activation)())
    return nn.Sequential(*layers)