Source code for bonni.model.mlp

from dataclasses import dataclass
from flax import linen as nn
import jax
from flax.linen import initializers

from bonni.model.utils import ActivationType, InitType, get_activation_fn, get_init_fn

class MLPLayer(nn.Module):
    out_channels: int
    activation_type: ActivationType
    norm_groups: int | None
    dropout_prob: float = 0.0
    bias_init: InitType = InitType.zeros
    skip_if_possible: bool = False
    
    def setup(self):
        bias_init_fn = get_init_fn(self.bias_init)
        self.fc1 = nn.Dense(
            features=self.out_channels, 
            bias_init=bias_init_fn,
            kernel_init=initializers.he_normal(),
        )
        self.norm = None
        if self.norm_groups is not None:
            self.norm = nn.GroupNorm(
                num_groups=self.norm_groups,
                epsilon=1e-5,  # Default epsilon value
            )
        self.activation = get_activation_fn(self.activation_type)
        self.dropout = None
        if self.dropout_prob > 0:
            self.dropout = nn.Dropout(
                rate=self.dropout_prob,
            )
    
    def __call__(
        self, 
        x: jax.Array, 
        deterministic: bool = False,
    ) -> jax.Array:
        # Apply first linear transformation
        y = self.fc1(x)
        # group norm
        if self.norm is not None:
            y = self.norm(y)
        # Apply activation
        y = self.activation(y)
        if self.skip_if_possible and x.shape == y.shape:
            y = x + y
        # Dropout
        if self.dropout is not None:
            y = self.dropout(y, deterministic=deterministic)
        return y




@dataclass(frozen=True, kw_only=True)
class MLPModelConfig:
    """
    Configuration object for a Multi-Layer Perceptron (MLP) model.

    This dataclass defines the structural and hyperparameter settings for an MLP,
    including layer dimensions, normalization, dropout, and activation strategies.
    It is frozen (immutable) and requires keyword arguments for initialization.

    Attributes:
        num_layer (int): The total number of linear layers in the MLP.
        out_channels (int): The dimensionality of the output features.
        hidden_channels (int | None): The dimensionality of the hidden layers. 
            If None, this is typically inferred from the input or output channels 
            depending on the implementation. Defaults to None.
        norm_groups (int | None): The number of groups to use for Group Normalization 
            in the hidden layers. If None, normalization is skipped. Defaults to None.
        last_norm_groups (int | None): The number of groups for Group Normalization 
            applied to the final layer. If None, no normalization is applied to the 
            output. Defaults to None.
        dropout_prob (float): The dropout probability applied after hidden layers. 
            Must be between 0.0 and 1.0. Defaults to 0.0.
        last_dropout_prob (float): The dropout probability applied after the final layer. 
            Defaults to 0.0.
        activation_type (ActivationType): The activation function used after hidden layers. 
            Defaults to ActivationType.gelu.
        different_last_activation (ActivationType | None): The activation function 
            used after the final layer. If set to `ActivationType.identity`, the output 
            is linear. If None, the model typically uses the same activation as 
            `activation_type`. Defaults to ActivationType.identity.
        bias_init (InitType): The initialization strategy for the layer biases 
            (e.g., zeros, uniform). Defaults to InitType.zeros.
        skip_if_possible (bool): If True, adds residual connections (skip connections) 
            around layers where the input and output dimensions are identical. 
            Defaults to True.
    """
    num_layer: int
    out_channels: int
    hidden_channels: int | None = None
    norm_groups: int | None = None
    last_norm_groups: int | None = None
    dropout_prob: float = 0.0
    last_dropout_prob: float = 0.0
    activation_type: ActivationType = ActivationType.gelu
    different_last_activation: ActivationType | None = ActivationType.identity  # if none, use same activation
    bias_init: InitType = InitType.zeros
    skip_if_possible: bool = True



class MLP(nn.Module):
    cfg: MLPModelConfig
    
    def setup(self):
        if self.cfg.num_layer > 1:
            assert self.cfg.hidden_channels is not None, "need hidden dim with >1 layer"
            
        layers = []
        for idx in range(self.cfg.num_layer):
            # select activation
            cur_activ = (
                self.cfg.activation_type 
                if idx != self.cfg.num_layer-1 or self.cfg.different_last_activation is None 
                else self.cfg.different_last_activation
            )
            # select out channels
            cur_out_channels = self.cfg.out_channels
            if idx != self.cfg.num_layer-1:
                assert self.cfg.hidden_channels is not None
                cur_out_channels = self.cfg.hidden_channels
            # select dropout prob, norm_groups
            dropout_prob = self.cfg.dropout_prob if idx != self.cfg.num_layer-1 else self.cfg.last_dropout_prob
            norm_groups = self.cfg.norm_groups if idx != self.cfg.num_layer-1 else self.cfg.last_norm_groups
            cur_skip_if_possible = self.cfg.skip_if_possible if idx != self.cfg.num_layer-1 else False
            # build current layer
            cur_layer = MLPLayer(
                out_channels=cur_out_channels,
                activation_type=cur_activ,
                norm_groups=norm_groups,
                dropout_prob=dropout_prob,
                bias_init=self.cfg.bias_init,
                skip_if_possible=cur_skip_if_possible,
            )
            layers.append(cur_layer)
        self.layers = layers
    
    def __call__(
        self,
        x: jax.Array,
        deterministic: bool = False,
    ) -> jax.Array:
        for layer in self.layers:
            x = layer(x, deterministic=deterministic)
        return x