Source code for bonni.model.utils

from enum import Enum
from typing import Callable
import jax
import jax.numpy as jnp
from flax import linen as nn


class SkipConnectionType(Enum):
    linear = "linear"
    identity = "identity"


def get_skip_connection(
    in_channels: int, 
    out_channels: int, 
    skip_type: SkipConnectionType,
):
    """Create a skip connection module."""
    if skip_type == SkipConnectionType.linear:
        return _LinearSkipConnection(in_channels, out_channels)
    if skip_type == SkipConnectionType.identity:
        if in_channels == out_channels:
            return lambda x: x
        else:
            return _LinearSkipConnection(in_channels, out_channels)
    raise ValueError(f"Unsupported skip connection type: {skip_type}")


class _LinearSkipConnection(nn.Module):
    """A linear skip connection implemented as a flax.linen Module."""
    in_features: int
    out_features: int
    
    def setup(self):
        self.linear = nn.Dense(
            features=self.out_features,
            use_bias=False,
            kernel_init=nn.initializers.lecun_normal(),
        )
    
    def __call__(self, x):
        return self.linear(x)




class ActivationType(Enum):
    """
    Enumeration of supported activation functions for neural network layers.

    These values are used to configure the non-linearity applied after linear 
    transformations in the model configuration.

    Attributes:
        identity: Applies no activation (f(x) = x). typically used for the final 
            output layer to produce unbounded linear predictions.
        gelu: Gaussian Error Linear Unit. A smooth approximation of ReLU often 
            used in Transformer architectures and modern MLPs.
        relu: Rectified Linear Unit (f(x) = max(0, x)). A standard non-linear 
            activation that introduces sparsity.
        leaky_relu: Leaky Rectified Linear Unit. Similar to ReLU but allows a 
            small, non-zero gradient when the unit is not active.
        sigmoid: Sigmoid function. Squashes values to the range [0, 1], often 
            used for binary classification probabilities.
        tanh: Hyperbolic Tangent. Squashes values to the range [-1, 1].
    """
    identity = "identity"
    gelu = "gelu"
    relu = "relu"
    leaky_relu = "leaky_relu"
    sigmoid = "sigmoid"
    tanh = "tanh"

    

def get_activation_fn(
    activation_type: ActivationType,
) -> Callable[[jax.Array], jax.Array]:
    if activation_type == ActivationType.identity:
        return lambda x: x
    if activation_type == ActivationType.gelu:
        return jax.nn.gelu
    if activation_type == ActivationType.relu:
        return jax.nn.relu
    if activation_type == ActivationType.leaky_relu:
        return jax.nn.leaky_relu
    if activation_type == ActivationType.sigmoid:
        return jax.nn.sigmoid
    if activation_type == ActivationType.tanh:
        return jax.nn.tanh
    raise ValueError(f"Invalid activation type: {activation_type}")




class InitType(Enum):
    """
    Enumeration of initialization strategies for model parameters.

    These values define how weights or biases are initialized before training begins.
    Used primarily for `bias_init` in the model configuration.

    Attributes:
        zeros: Initializes parameters to exactly 0. This is the standard practice 
            for bias terms in most neural network layers.
        ones: Initializes parameters to exactly 1.
        uniform: Initializes parameters with values drawn from a uniform distribution.
            The range is typically determined by the specific layer implementation.
        normal: Initializes parameters with values drawn from a normal (Gaussian) 
            distribution.
    """
    zeros = "zeros"
    ones = "ones"
    uniform = "uniform"
    normal = "normal"



def get_init_fn(
    init_type: InitType,
):
    if init_type == InitType.zeros:
        return nn.initializers.zeros
    if init_type == InitType.ones:
        return nn.initializers.ones
    if init_type == InitType.uniform:
        scale = 0.01
        def init(
            key, 
            shape, 
            dtype=jnp.float64, 
            out_sharding=None
        ) -> jax.Array:
            return jax.random.uniform(key, shape, dtype=dtype, out_sharding=out_sharding) * jnp.array(scale, dtype)
        return jax.tree_util.Partial(init)
    if init_type == InitType.normal:
        scale = 0.1
        def init(
            key, 
            shape, 
            dtype=jnp.float64, 
            out_sharding=None
        ) -> jax.Array:
            return jax.random.normal(key, shape, dtype=dtype, out_sharding=out_sharding) * jnp.array(scale, dtype)
        return jax.tree_util.Partial(init)
    raise ValueError(f"Invalid init type: {init_type}")