Bayesian Layers

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Applies linear transformation using deterministic means in eval mode or LRT sampling in train mode.

    Args:
        x: Input tensor of shape (..., in_features).

    Returns:
        torch.Tensor: Output tensor of shape (..., out_features).
    """
    if not self.training:
        return F.linear(x, self.w_mu, self.b_mu)

    w_sigma = F.softplus(self.w_rho)
    b_sigma = F.softplus(self.b_rho)

    gamma = F.linear(x, self.w_mu)
    delta = F.linear(x.pow(2), w_sigma.pow(2)) + b_sigma.pow(2)

    eps = torch.randn_like(gamma)
    return gamma + eps * torch.sqrt(delta + 1e-8) + self.b_mu

def reset_parameters(self) -> None:
    """Initializes layer weights and bias parameters."""
    if self.weight_init == "kaiming":
        init.kaiming_normal_(self.w_mu, nonlinearity="relu")
    elif self.weight_init == "xavier":
        init.xavier_normal_(self.w_mu)
    else:
        init.kaiming_uniform_(self.w_mu, a=math.sqrt(5))

    init.zeros_(self.b_mu)

    rho_init_val = math.log(math.exp(self.init_sigma) - 1.0)
    self.w_rho.data.fill_(rho_init_val)
    self.b_rho.data.fill_(rho_init_val)

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Executes forward pass using deterministic means in eval mode or LRT sampling in train mode.

    Args:
        x: Input tensor of shape (batch_size, in_channels, height, width).

    Returns:
        torch.Tensor: Convolved output tensor of shape (batch_size, out_channels, out_height, out_width).
    """
    if not self.training:
        return F.conv2d(
            x,
            self.w_mu,
            self.b_mu,
            self.stride,
            self.padding,
            self.dilation,
            self.groups,
        )

    w_sigma = F.softplus(self.w_rho)
    b_sigma = F.softplus(self.b_rho)

    conv_mu = F.conv2d(
        x,
        self.w_mu,
        bias=None,
        stride=self.stride,
        padding=self.padding,
        dilation=self.dilation,
        groups=self.groups,
    )

    conv_var = F.conv2d(
        x.pow(2),
        w_sigma.pow(2),
        bias=None,
        stride=self.stride,
        padding=self.padding,
        dilation=self.dilation,
        groups=self.groups,
    )

    conv_var = conv_var + b_sigma.pow(2).view(1, -1, 1, 1)

    eps = torch.randn_like(conv_mu)
    out = conv_mu + eps * torch.sqrt(conv_var + 1e-8)

    return out + self.b_mu.view(1, -1, 1, 1)

def reset_parameters(self) -> None:
    """Initializes layer weights and bias parameters."""
    init.kaiming_normal_(self.w_mu, mode="fan_in", nonlinearity="relu")
    init.zeros_(self.b_mu)

    rho_init = math.log(math.exp(self.init_sigma) - 1.0)
    self.w_rho.data.fill_(rho_init)
    self.b_rho.data.fill_(rho_init)

def apply_pruning(self, threshold: float = 0.83) -> float:
    """
    Applies pruning in-place: zeros out the means and minimizes the variance
    for weights that fall below the SNR threshold.

    Returns:
        float: Sparsity of the layer (percentage of pruned weights, 0.0 to 1.0).
    """
    masks = self.get_pruning_masks(threshold)

    total_weights = 0
    pruned_weights = 0

    with torch.no_grad():
        for name, param in self.named_parameters():
            if name.endswith("_mu"):
                mask = masks[name]

                total_weights += mask.numel()
                pruned_weights += (mask == 0.0).sum().item()
                param.data *= mask

                rho_name = name.replace("_mu", "_rho")
                if hasattr(self, rho_name):
                    rho = getattr(self, rho_name)
                    rho.data = torch.where(
                        mask.bool(),
                        rho.data,
                        torch.tensor(-1e8, device=rho.device, dtype=rho.dtype),
                    )

    sparsity = pruned_weights / total_weights if total_weights > 0 else 0.0
    return sparsity

def get_pruning_masks(self, threshold: float = 0.83) -> dict:
    """Returns binary masks for parameters satisfying the SNR threshold.

    Implements Graves' pruning heuristic where weights with low
    Signal-to-Noise Ratio are considered redundant and can be removed.

    Args:
        threshold (float, optional): The SNR threshold (|mu|/sigma).
            Defaults to 0.83, the "safe" threshold suggested by Graves.

    Returns:
        dict[str, torch.Tensor]: A dictionary mapping parameter names to
            binary masks (1.0 for keeping, 0.0 for pruning).
    """
    snr_dict = self._get_snr_dict()
    return {name: (val > threshold).float() for name, val in snr_dict.items()}

def kl_divergence(self) -> torch.Tensor:
    """
    Computes KL-divergence KL(q || p) for all bayesian weights of the layer.
    p(w) = N(0, prior_sigma^2)
    q(w) = N(mu, sigma^2), where sigma = softplus(rho)
    """
    total_kl = 0.0

    for name, param in self.named_parameters():
        if name.endswith("_mu"):
            rho_name = name.replace("_mu", "_rho")

            if hasattr(self, rho_name):
                mu = param
                rho = getattr(self, rho_name)

                total_kl += self._compute_kl_for_param(mu, rho)

    return total_kl