3. Neural Network Layers

class BayesianModule(nnx.Module, ABC):
    """
    Abstract base for Bayesian-aware modules in JAX.
    Provides mechanisms to track if a module is Bayesian and control
    parameter updates through freezing/unfreezing.

    Notes:
        All derived classes must implement `freeze` and `kl_cost` to
        handle parameter management and compute the KL divergence cost.
    """

    def __init__(self, **kwargs: Any) -> None:
        """
        Initialize the Bayesian module with default flags.
        Sets `frozen` to False and `is_bayesian` to True.

        Args:
            **kwargs: Extra arguments passed to the base class.

        Returns:
            None.
        """

        super().__init__(**kwargs)

        self.frozen: bool = False
        self.is_bayesian: bool = True

    @abstractmethod
    def freeze(self) -> None:
        """
        Freeze the module's parameters to stop gradient computation.
        If weights or biases are not sampled yet, they are sampled first.
        Once frozen, parameters are not resampled or updated.

        Returns:
            None.

        Notes:
            Must be implemented by all subclasses.
        """

    def unfreeze(self) -> None:
        """
        Unfreeze the module by setting its `frozen` flag to False.
        Allows parameters to be sampled and updated again.

        Returns:
            None.
        """

        self.frozen = False

    @abstractmethod
    def kl_cost(self) -> tuple[jax.Array, int]:
        """
        Compute the KL divergence cost for all Bayesian parameters.

        Returns:
            tuple[jax.Array, int]: A tuple containing the KL divergence
                cost and the total number of parameters in the layer.

        Notes:
            Must be implemented by all subclasses.
        """

class Conv1d(BayesianModule):
    """
    Bayesian 1D convolutional layer with optional weight and bias priors.
    Behaves like a standard Conv1d but treats weights and bias as random
    variables sampled from specified distributions. Parameters become fixed
    when the layer is frozen.
    """

    bias_distribution: Optional[GaussianDistribution] = None
    bias: Optional[nnx.Param] = None

    def __init__(
        self,
        input_channels: int,
        output_channels: int,
        kernel_size: int,
        stride: int = 1,
        padding: int = 0,
        dilation: int = 1,
        groups: int = 1,
        weights_distribution: Optional[GaussianDistribution] = None,
        bias_distribution: Optional[GaussianDistribution] = None,
        use_bias: bool = True,
        rngs: Rngs = nnx.Rngs(0),
        **kwargs: Any,
    ) -> None:
        """
        Initializes a Bayesian 1D convolutional layer.

        Args:
            input_channels: Number of input feature channels.
            output_channels: Number of output feature channels.
            kernel_size: Size of the convolution kernel.
            stride: Stride of the convolution operation.
            padding: Amount of zero-padding on both sides.
            dilation: Spacing between kernel elements.
            groups: Number of blocked connections between input and output.
            weights_distribution: Distribution to initialize weights.
            bias_distribution: Distribution to initialize bias.
            use_bias: Whether to include a bias term.
            rngs: Random number generators for reproducibility.
            **kwargs: Extra arguments passed to the base class.

        Returns:
            None.

        Notes:
            Gaussian distributions are used by default if none are
            provided.
        """

        super().__init__(**kwargs)

        self.input_channels = input_channels
        self.output_channels = output_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups
        self.use_bias = use_bias
        self.rngs = rngs

        # Set weights prior
        if weights_distribution is None:
            self.weights_distribution = GaussianDistribution(
                shape=(
                    self.output_channels,
                    self.input_channels // self.groups,
                    self.kernel_size,
                ),
                rngs=self.rngs,
            )
        else:
            self.weights_distribution = weights_distribution

        # Set bias prior
        if self.use_bias:
            if bias_distribution is None:
                self.bias_distribution = GaussianDistribution(
                    shape=(self.output_channels,),
                    rngs=self.rngs,
                )
            else:
                self.bias_distribution = bias_distribution
        else:
            self.bias_distribution = None

        # Sample initial weights
        self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))

        # Sample initial bias only if using bias
        if self.use_bias and self.bias_distribution is not None:
            self.bias = nnx.Param(self.bias_distribution.sample(self.rngs))
        else:
            self.bias = None

    def freeze(self) -> None:
        """
        Freeze the module's parameters to stop gradient computation.
        If weights or biases are not sampled yet, they are sampled first.
        Once frozen, parameters are not resampled or updated.

        Returns:
            None.
        """

        # Set indicator
        self.frozen = True

        # Sample weights if they are undefined
        if self.weights is None:
            self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))

        # Sample bias if they are undefined and bias is used
        if self.use_bias and self.bias is None and self.bias_distribution is not None:
            self.bias = nnx.Param(self.bias_distribution.sample(self.rngs))

        # Stop gradient computation
        self.weights = jax.lax.stop_gradient(self.weights)
        if self.use_bias:
            self.bias = jax.lax.stop_gradient(self.bias)

    def kl_cost(self) -> tuple[jax.Array, int]:
        """
        Compute the KL divergence cost for all Bayesian parameters.

        Returns:
            tuple[jax.Array, int]: A tuple containing the KL divergence
                cost and the total number of parameters in the layer.
        """

        # Compute log probs for weights
        log_probs: jax.Array = self.weights_distribution.log_prob(
            jnp.asarray(self.weights)
        )

        # Add bias log probs only if using bias
        if (
            self.use_bias
            and self.bias is not None
            and self.bias_distribution is not None
        ):
            log_probs += self.bias_distribution.log_prob(jnp.asarray(self.bias))

        # Compute number of parameters
        num_params: int = self.weights_distribution.num_params
        if self.use_bias and self.bias_distribution is not None:
            num_params += self.bias_distribution.num_params

        return log_probs, num_params

    def __call__(self, inputs: jax.Array) -> jax.Array:
        """
        Performs a forward pass through the Bayesian Convolution 1D
        layer. If the layer is not frozen, it samples weights and bias
        from their respective distributions. If the layer is frozen
        and the weights or bias are not initialized, it also performs
        sampling.

        Args:
            inputs: Input tensor to the layer with shape
                (batch, channels, length).

        Returns:
            Output array after convolution with optional bias added.

        Raises:
            ValueError: If the layer is frozen but weights or bias are
                undefined.
        """

        # Sample if model not frozen
        if not self.frozen:
            # Sample weights
            self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))

            # Sample bias only if using bias
            if self.use_bias and self.bias_distribution is not None:
                self.bias = nnx.Param(self.bias_distribution.sample(self.rngs))
        elif self.weights is None or (self.use_bias and self.bias is None):
            raise ValueError(
                "Module has been frozen with undefined weights and/or bias."
            )

        # Compute outputs
        outputs = jax.lax.conv_general_dilated(
            lhs=inputs,
            rhs=jnp.asarray(self.weights),
            window_strides=[self.stride],
            padding=[(self.padding, self.padding)],
            lhs_dilation=[1],
            rhs_dilation=[self.dilation],
            dimension_numbers=(
                "NCH",  # Input
                "OIH",  # Kernel
                "NCH",  # Output
            ),
            feature_group_count=self.groups,
        )

        # Add bias only if using bias
        if self.use_bias and self.bias is not None:
            outputs += jnp.reshape(
                a=jnp.asarray(self.bias), shape=(1, self.output_channels, 1)
            )

        return outputs

class Conv2d(BayesianModule):
    """
    Bayesian 2D convolutional layer with optional weight and bias priors.
    Behaves like a standard Conv2d but treats weights and bias as random
    variables sampled from specified distributions. Parameters become fixed
    when the layer is frozen.
    """

    bias_distribution: Optional[GaussianDistribution] = None
    bias: Optional[nnx.Param] = None

    def __init__(
        self,
        input_channels: int,
        output_channels: int,
        kernel_size: int | tuple[int, int],
        stride: tuple[int, int] = (1, 1),
        padding: tuple[int, int] = (0, 0),
        dilation: tuple[int, int] = (1, 1),
        groups: int = 1,
        weights_distribution: Optional[GaussianDistribution] = None,
        bias_distribution: Optional[GaussianDistribution] = None,
        use_bias: bool = True,
        rngs: Rngs = nnx.Rngs(0),
        **kwargs: Any,
    ) -> None:
        """
        Initializes a Bayesian 2D convolutional layer.

        Args:
            input_channels: Number of input feature channels.
            output_channels: Number of output feature channels.
            kernel_size: Convolution kernel size. Int is converted to tuple.
            stride: Stride of the convolution operation.
            padding: Tuple specifying zero-padding for height and width.
            dilation: Spacing between kernel elements.
            groups: Number of blocked connections between input and output.
            weights_distribution: Distribution to initialize weights.
            bias_distribution: Distribution to initialize bias.
            use_bias: Whether to include a bias term.
            rngs: Random number generators for reproducibility.
            **kwargs: Extra arguments passed to the base class.

        Returns:
            None.

        Notes:
            Gaussian distributions are used by default if none are
            provided.
        """

        super().__init__(**kwargs)

        self.input_channels = input_channels
        self.output_channels = output_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups
        self.use_bias = use_bias
        self.rngs = rngs

        # Set weights distribution
        if weights_distribution is None:
            # Extend kernel if we only have 1 value
            if isinstance(self.kernel_size, int):
                self.kernel_size = (self.kernel_size, self.kernel_size)

            self.weights_distribution: GaussianDistribution = GaussianDistribution(
                shape=(
                    self.output_channels,
                    self.input_channels // self.groups,
                    *self.kernel_size,
                ),
                rngs=self.rngs,
            )
        else:
            self.weights_distribution = weights_distribution

        # Set bias prior
        if self.use_bias:
            if bias_distribution is None:
                # Define weights distribution
                self.bias_distribution = GaussianDistribution(
                    shape=(self.output_channels,),
                    rngs=self.rngs,
                )
            else:
                self.bias_distribution = bias_distribution
        else:
            self.bias_distribution = None

        # Sample initial weights
        self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))

        # Sample initial bias only if using bias
        if self.use_bias and self.bias_distribution is not None:
            self.bias = nnx.Param(self.bias_distribution.sample(self.rngs))
        else:
            self.bias = None

    def freeze(self) -> None:
        """
        Freeze the module's parameters to stop gradient computation.
        If weights or biases are not sampled yet, they are sampled first.
        Once frozen, parameters are not resampled or updated.

        Returns:
            None.
        """

        # Set indicator
        self.frozen = True

        # Sample weights if they are undefined
        if self.weights is None:
            self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))

        # Sample bias if they are undefined and bias is used
        if self.use_bias and self.bias is None and self.bias_distribution is not None:
            self.bias = nnx.Param(self.bias_distribution.sample(self.rngs))

        # Stop gradient computation
        self.weights = jax.lax.stop_gradient(self.weights)
        if self.use_bias:
            self.bias = jax.lax.stop_gradient(self.bias)

    def kl_cost(self) -> tuple[jax.Array, int]:
        """
        Compute the KL divergence cost for all Bayesian parameters.

        Returns:
            tuple[jax.Array, int]: A tuple containing the KL divergence
                cost and the total number of parameters in the layer.
        """

        # Compute log probs for weights
        log_probs: jax.Array = self.weights_distribution.log_prob(
            jnp.asarray(self.weights)
        )

        # Add bias log probs only if using bias
        if (
            self.use_bias
            and self.bias is not None
            and self.bias_distribution is not None
        ):
            log_probs += self.bias_distribution.log_prob(jnp.asarray(self.bias))

        # Compute number of parameters
        num_params: int = self.weights_distribution.num_params
        if self.use_bias and self.bias_distribution is not None:
            num_params += self.bias_distribution.num_params

        return log_probs, num_params

    def __call__(self, inputs: jax.Array) -> jax.Array:
        """
        Performs a forward pass through the Bayesian Convolution 2D
        layer. If the layer is not frozen, it samples weights and bias
        from their respective distributions. If the layer is frozen
        and the weights or bias are not initialized, it also performs
        sampling.

        Args:
            inputs: Input array with shape (batch, channels, height,
                width).

        Returns:
            Output array after convolution with optional bias addition.

        Raises:
            ValueError: If the layer is frozen but weights or bias are
                undefined.
        """

        # Sample if model not frozen
        if not self.frozen:
            # Sample weights
            self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))

            # Sample bias only if using bias
            if self.use_bias and self.bias_distribution is not None:
                self.bias = nnx.Param(self.bias_distribution.sample(self.rngs))
        elif self.weights is None or (self.use_bias and self.bias is None):
            raise ValueError(
                "Module has been frozen with undefined weights and/or bias."
            )

        # Compute ouputs
        outputs = jax.lax.conv_general_dilated(
            lhs=inputs,
            rhs=jnp.asarray(self.weights),
            window_strides=self.stride,
            padding=[self.padding, self.padding],
            lhs_dilation=[1, 1],
            rhs_dilation=self.dilation,
            dimension_numbers=(
                "NCHW",  # Input
                "OIHW",  # Kernel
                "NCHW",  # Output
            ),
            feature_group_count=self.groups,
        )

        # Add bias only if using bias
        if self.use_bias and self.bias is not None:
            outputs += jnp.reshape(
                a=jnp.asarray(self.bias), shape=(1, self.output_channels, 1, 1)
            )

        return outputs

class Embedding(BayesianModule):
    """
    Bayesian embedding layer with optional padding and max-norm constraints.
    Each embedding vector is sampled from a specified weight distribution.
    If the layer is frozen, embeddings are fixed and gradients are stopped.
    """

    def __init__(
        self,
        num_embeddings: int,
        embeddings_dim: int,
        padding_idx: Optional[int] = None,
        max_norm: Optional[float] = None,
        norm_type: float = 2.0,
        scale_grad_by_freq: bool = False,
        weights_distribution: Optional[GaussianDistribution] = None,
        rngs: Rngs = nnx.Rngs(0),
        **kwargs: Any,
    ) -> None:
        """
        Initialize a Bayesian embedding layer with optional constraints.
        Sets up the embedding weight distribution and samples initial values.

        Args:
            num_embeddings: Size of the embedding dictionary.
            embeddings_dim: Dimension of each embedding vector.
            padding_idx: Index whose embeddings are ignored in gradient.
            max_norm: Maximum norm for each embedding vector.
            norm_type: p value for the p-norm in max_norm option.
            scale_grad_by_freq: Scale gradients by inverse word frequency.
            weights_distribution: Distribution to initialize embeddings.
            rngs: Random number generators for reproducibility.
            **kwargs: Extra arguments passed to the base class.

        Returns:
            None.

        Notes:
            Gaussian distributions are used by default if none are
            provided.
        """

        super().__init__(**kwargs)

        self.num_embeddings = num_embeddings
        self.embeddings_dim = embeddings_dim
        self.padding_idx = padding_idx
        self.max_norm = max_norm
        self.norm_type = norm_type
        self.scale_grad_by_freq = scale_grad_by_freq
        self.rngs = rngs

        # Set weights distribution
        if weights_distribution is None:
            self.weights_distribution = GaussianDistribution(
                (self.num_embeddings, self.embeddings_dim)
            )
        else:
            self.weights_distribution = weights_distribution

        # Sample initial weights
        self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))

    def freeze(self) -> None:
        """
        Freeze the module's parameters to stop gradient computation.
        If weights or biases are not sampled yet, they are sampled first.
        Once frozen, parameters are not resampled or updated.

        Returns:
            None.
        """

        # Set indicator
        self.frozen = True

        # Sample weights if they are undefined
        if self.weights is None:
            self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))

        # Stop gradient computation
        self.weights = jax.lax.stop_gradient(self.weights)

    def kl_cost(self) -> tuple[jax.Array, int]:
        """
        Compute the KL divergence cost for all Bayesian parameters.

        Returns:
            tuple[jax.Array, int]: A tuple containing the KL divergence
                cost and the total number of parameters in the layer.
        """

        # Compute log probs for weights
        log_probs: jax.Array = self.weights_distribution.log_prob(
            jnp.asarray(self.weights)
        )

        # get number of parameters
        num_params: int = self.weights_distribution.num_params

        return log_probs, num_params

    def __call__(self, inputs: jax.Array) -> jax.Array:
        """
        Perform a forward pass using current embedding weights.

        Args:
            inputs: Array of indices into the embedding matrix.

        Returns:
            Array of shape [*, embeddings_dim] containing the embedding
            vectors corresponding to the input indices.

        Raises:
            ValueError: If the layer is frozen but weights are
                undefined.

        Notes:
            Embeddings at padding_idx are zeroed out, and vectors exceeding
            max_norm are renormalized if specified.
        """

        # Sample if model not frozen
        if not self.frozen:
            # Sample weights
            self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))
        elif self.weights is None:
            raise ValueError("Module has been frozen with undefined weights.")

        # Perform embedding lookup
        outputs = self.weights.value[inputs]

        # Apply padding_idx
        if self.padding_idx is not None:
            # Create mask for padding indices
            mask = inputs == self.padding_idx
            # Zero out embeddings for padding indices
            outputs = jnp.where(mask[..., None], 0.0, outputs)

        # Apply max_norm
        if self.max_norm is not None:
            norms = jnp.linalg.norm(outputs, axis=-1, ord=self.norm_type, keepdims=True)
            # Normalize vectors that exceed max_norm
            scale = jnp.minimum(1.0, self.max_norm / (norms + 1e-8))
            outputs = outputs * scale

        return outputs

class Linear(BayesianModule):
    """
    Bayesian linear (fully connected) layer with optional weight and bias
    priors. Functions like a standard linear layer but treats weights and
    bias as probabilistic variables. Freezing the layer fixes parameters
    and stops gradient computation.
    """

    bias_distribution: Optional[GaussianDistribution] = None
    bias: Optional[nnx.Param] = None

    def __init__(
        self,
        input_size: int,
        output_size: int,
        weights_distribution: Optional[GaussianDistribution] = None,
        bias_distribution: Optional[GaussianDistribution] = None,
        use_bias: bool = True,
        precision: PrecisionLike = None,
        dot_general: DotGeneralT = lax.dot_general,
        rngs: Rngs = nnx.Rngs(0),
        **kwargs: Any,
    ) -> None:
        """
        Initialize a Bayesian linear layer with optional priors for weights
        and bias. Samples initial parameter values from the specified
        distributions.

        Args:
            input_size: Number of input features.
            output_size: Number of output features.
            weights_distribution: Distribution for weights.
            bias_distribution: Distribution for bias.
            use_bias: Whether to include a bias term.
            precision: Precision for dot product computations.
            dot_general: Function for generalized dot products.
            rngs: Random number generators for reproducibility.
            **kwargs: Extra arguments passed to the base class.

        Returns:
            None.

        Notes:
            Gaussian distributions are used by default if none are
            provided.
        """

        super().__init__(**kwargs)

        self.input_size = input_size
        self.output_size = output_size
        self.use_bias = use_bias
        self.precision = precision
        self.dot_general = dot_general
        self.rngs = rngs

        # Set weights prior
        if weights_distribution is None:
            self.weights_distribution = GaussianDistribution(
                shape=(self.output_size, self.input_size), rngs=self.rngs
            )
        else:
            self.weights_distribution = weights_distribution

        # Set bias prior
        if self.use_bias:
            if bias_distribution is None:
                self.bias_distribution = GaussianDistribution(
                    shape=(self.output_size,), rngs=self.rngs
                )
            else:
                self.bias_distribution = bias_distribution
        else:
            self.bias_distribution = None

        # Sample initial weights
        self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))

        # Sample initial bias only if using bias
        if self.use_bias and self.bias_distribution is not None:
            self.bias = nnx.Param(self.bias_distribution.sample(self.rngs))
        else:
            self.bias = None

    def freeze(self) -> None:
        """
        Freeze the module's parameters to stop gradient computation.
        If weights or biases are not sampled yet, they are sampled first.
        Once frozen, parameters are not resampled or updated.

        Returns:
            None.
        """

        # Set indicator
        self.frozen = True

        # Sample weights if they are undefined
        if self.weights is None:
            self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))

        # Sample bias if they are undefined and bias is used
        if self.use_bias and self.bias is None and self.bias_distribution is not None:
            self.bias = nnx.Param(self.bias_distribution.sample(self.rngs))

        # Stop gradient computation
        self.weights = jax.lax.stop_gradient(self.weights)
        if self.use_bias:
            self.bias = jax.lax.stop_gradient(self.bias)

    def kl_cost(self) -> tuple[jax.Array, int]:
        """
        Compute the KL divergence cost for all Bayesian parameters.

        Returns:
            tuple[jax.Array, int]: A tuple containing the KL divergence
                cost and the total number of parameters in the layer.
        """

        # Compute log probs for weights
        log_probs: jax.Array = self.weights_distribution.log_prob(
            jnp.asarray(self.weights)
        )

        # Add bias log probs only if using bias
        if (
            self.use_bias
            and self.bias is not None
            and self.bias_distribution is not None
        ):
            log_probs += self.bias_distribution.log_prob(jnp.asarray(self.bias))

        # Compute number of parameters
        num_params: int = self.weights_distribution.num_params
        if self.use_bias and self.bias_distribution is not None:
            num_params += self.bias_distribution.num_params

        return log_probs, num_params

    def __call__(self, inputs: jax.Array) -> jax.Array:
        """
        Perform a forward pass using current weights and bias. Samples new
        parameters if the layer is not frozen.

        Args:
            inputs: Input array with shape [*, input_size].

        Returns:
            Output array with shape [*, output_size].

        Raises:
            ValueError: If the layer is frozen but weights or bias are
                undefined.
        """

        # Sample if model not frozen
        if not self.frozen:
            # Sample weights
            self.weights = nnx.Param(self.weights_distribution.sample(self.rngs))

            # Sample bias only if using bias
            if self.use_bias and self.bias_distribution is not None:
                self.bias = nnx.Param(self.bias_distribution.sample(self.rngs))
        elif self.weights is None or (self.use_bias and self.bias is None):
            raise ValueError(
                "Module has been frozen with undefined weights and/or bias."
            )

        # Compute outputs
        outputs = inputs @ self.weights.T

        # Add bias only if using bias
        if self.use_bias and self.bias is not None:
            outputs += jnp.reshape(
                jnp.asarray(self.bias), (1,) * (outputs.ndim - 1) + (-1,)
            )

        return outputs