kaira.models.components.MLPEncoder

Inheritance diagram for MLPEncoder

class kaira.models.components.MLPEncoder(in_features: int, out_features: int, hidden_dims: List[int] | None = None, activation: Module | None = None, output_activation: Module | None = None, *args: Any, **kwargs: Any)

Bases: BaseModel

Multi-Layer Perceptron (MLP) Encoder for communication systems.

This module implements a simple MLP-based encoder that maps input messages to encoded signals suitable for transmission over a communication channel.

Methods

__init__	Initialize the MLPEncoder.
forward	Forward pass of the MLPEncoder.

__init__(in_features: int, out_features: int, hidden_dims: List[int] | None = None, activation: Module | None = None, output_activation: Module | None = None, *args: Any, **kwargs: Any)

Initialize the MLPEncoder.

Parameters:

• in_features (int) – The dimensionality of the input messages.

• out_features (int) – The dimensionality of the output encoded signals.

• hidden_dims (List[int], optional) – Dimensions of hidden layers. If None, a single hidden layer with (in_features + out_features) // 2 units is used.

• activation (nn.Module, optional) – Activation function to use between layers. If None, ReLU is used.

• output_activation (nn.Module, optional) – Activation function to use at the output. If None, no activation is applied to the output.

• *args – Variable positional arguments passed to the base class.

• **kwargs – Variable keyword arguments passed to the base class.

forward(x: Tensor, *args: Any, **kwargs: Any) → Tensor

Forward pass of the MLPEncoder.

Parameters:

• x (torch.Tensor) – Input tensor of shape (batch_size, in_features).

• *args – Additional positional arguments (unused).

• **kwargs – Additional keyword arguments (unused).

Returns:

Output tensor of shape (batch_size, out_features).

Return type:

torch.Tensor