src.mlp

Classes

MLP(*args, **kwargs)

Multi-layer perceptron built from nn.Linear and activation layers.

class src.mlp.MLP(*args: Any, **kwargs: Any)

Bases: Module

Multi-layer perceptron built from nn.Linear and activation layers.

This module implements a feedforward neural network of the form:

\[f(x) = W_n \sigma( W_{n-1} \sigma(\dots \sigma(W_1 x + b_1)\dots ) + b_{n-1}) + b_n\]

Variables:

layers (nn.Sequential) – The sequential container holding all linear and activation layers.

Example

>>> mlp = MLP(input_dim=10, hidden_dims=[20, 30], output_dim=5, activation=nn.ReLU)
>>> x = torch.randn(4, 10)  # Batch of 4 samples
>>> output = mlp(x)
>>> print(output.shape)
torch.Size([4, 5])

Notes

• No activation is applied after the final (output) linear layer.

• Activation arguments should be activation classes (callables that return nn.Module instances), not instantiated modules.

Initialize the Multilayer Perceptron (MLP).

Parameters:

• input_dim (int) – Dimension of the input features.

• hidden_dims (int or Sequence[int]) – Size(s) of the hidden layer(s).

• output_dim (int) – Dimension of the output features.

• activation (Type[nn.Module] or Sequence[Type[nn.Module]], optional) – Activation function(s) to use. Defaults to nn.ReLU.

Notes

The number of hidden layers equals len(hidden_dims). If activation is a single class, the same activation is applied to all hidden layers.

forward(x: torch.Tensor) → torch.Tensor

Forward pass through the MLP.

The MLP is applied independently to the last dimension, so additional batch dimensions are preserved.

Parameters:

x (torch.Tensor) – Input tensor of shape (…, input_dim).

Returns:

Output tensor of shape (…, output_dim).

Return type:

torch.Tensor