2. Losses

class KLDivergenceLoss(nnx.Module):
    """
    Compute Kullback-Leibler divergence across Bayesian modules.
    This loss sums the KL divergence from all Bayesian layers in
    the model. It can be reduced by averaging and scaled by a
    weight factor.

    Notes:
        Assumes the model contains submodules derived from
        `BayesianModule`.
    """

    def __init__(
        self,
        reduction: Literal["mean"] = "mean",
        weight: float = 1.0,
        **kwargs: Any,
    ) -> None:
        """
        Initialize the KL divergence loss.

        Args:
            reduction: Method used to reduce the KL loss.
            weight: Scaling factor for the KL divergence.
            **kwargs: Extra arguments passed to the base class.

        Returns:
            None
        """

        super().__init__(**kwargs)

        self.reduction = reduction
        self.weight = weight

    def __call__(self, model: nnx.Module) -> jax.Array:
        """
        Compute KL divergence for all Bayesian modules in a model.

        Args:
            model: Model containing Bayesian submodules.

        Returns:
            jax.Array: Weighted KL divergence loss.

        Notes:
            The KL loss is averaged over the number of parameters
            and scaled by the `weight` attribute.
        """

        # Init kl cost and params
        kl_global_cost: jax.Array = jnp.array(0.0)
        num_params_global: int = 0

        # Iter over modules
        for _, module in model.iter_modules():
            if isinstance(module, BayesianModule):
                kl_cost, num_params = module.kl_cost()
                kl_global_cost += kl_cost
                num_params_global += num_params

        # Average by the number of parameters
        kl_global_cost /= num_params
        kl_global_cost *= self.weight

        return kl_global_cost

class ELBOLoss(nnx.Module):
    """
    Compute the Evidence Lower Bound (ELBO) loss for Bayesian
    networks. Combines a reconstruction loss with a KL divergence
    term. Monte Carlo sampling can estimate the expected
    reconstruction loss over stochastic layers.

    Notes:
        The KL term is weighted by `kl_weight`. The model is
        assumed to contain Bayesian layers compatible with
        `KLDivergenceLoss`.
    """

    def __init__(
        self,
        loss_function: Callable[[jax.Array, jax.Array], jax.Array],
        num_samples: int = 1,
        kl_weight: float = 1e-3,
        **kwargs: Any,
    ) -> None:
        """
        Initialize the ELBO loss with reconstruction and KL
        components.

        Args:
            loss_function: Function to compute reconstruction
                loss.
            num_samples: Number of Monte Carlo samples used for
                estimation.
            kl_weight: Weight applied to the KL divergence term.
            **kwargs: Extra arguments passed to the base class.

        Returns:
            None
        """

        super().__init__(**kwargs)

        self.loss_function = loss_function
        self.num_samples = num_samples
        self.kl_weight = kl_weight
        self.kl_loss = KLDivergenceLoss(weight=kl_weight)

    def __call__(
        self, outputs: jax.Array, targets: jax.Array, model: nnx.Module
    ) -> jax.Array:
        """
        Compute the ELBO loss with Monte Carlo sampling and KL
        regularization.

        Args:
            outputs: Predictions generated by the model.
            targets: Ground truth values for training.
            model: Model containing Bayesian layers.

        Returns:
            jax.Array: Scalar ELBO loss averaged over samples.

        Notes:
            The loss is averaged over `num_samples` Monte Carlo
            draws.
        """

        loss_value: jax.Array = jnp.array(0.0)
        for _ in range(self.num_samples):
            loss_value += self.loss_function(outputs, targets) + self.kl_loss(model)

        loss_value /= self.num_samples

        return loss_value