Posterior Approximation Methods

def build_ensemble(self, n_members: int = 10, temperature: float = 1.0) -> Ensemble:
    """Builds an Ensemble module sampled from the Laplace posterior.

    Args:
        n_members: Number of ensemble members to draw. Defaults to 10.
        temperature: Sampling temperature for posterior weights. Defaults to 1.0.

    Returns:
        Ensemble: Ensemble instance wrapping the sampled models.
    """
    return Ensemble.from_posterior(
        self, n_members=n_members, temperature=temperature
    )

def compute_curvature(
    self,
    train_loader: DataLoader,
    num_samples: int = 1000,
) -> None:
    """Estimates the Kronecker factors of the Hessian using training data.

    Args:
        train_loader: DataLoader yielding training data and targets.
        num_samples: Maximum number of samples to process for curvature estimation. Defaults to 1000.
    """
    self.dataset_size = len(train_loader.dataset)

    if self.verbose:
        print("Registering hooks...")
    self._register_hooks()

    try:
        if self.verbose:
            print("Estimating Kronecker factors...")
        self._estimate_kronecker_factors(train_loader, num_samples)
    finally:
        if self.verbose:
            print("Removing hooks...")
        self._remove_hooks()

    self.is_fitted = True

    if self.verbose:
        print("Curvature computation completed!")

def sample_models(
    self, n_models: int = 10, temperature: float = 1.0
) -> list[nn.Module]:
    """Samples model parameters from the approximated Gaussian posterior.

    Args:
        n_models: Number of models to sample. Defaults to 10.
        temperature: Sampling temperature scaling factor. Defaults to 1.0.

    Returns:
        list[nn.Module]: Sampled model instances in eval mode.

    Raises:
        RuntimeError: If curvature has not been computed prior to sampling.
    """
    if not self.is_fitted:
        raise RuntimeError(
            "Laplace curvature not computed. Call compute_curvature() first."
        )

    samples = []
    modules = dict(self.model.named_modules())

    for _ in range(n_models):
        sampled_state = copy.deepcopy(self.model.state_dict())

        for name, factors in self.sampling_factors.items():
            module = modules[name]

            mean_weight = module.weight.detach()
            l_q = factors["L_U"].to(device=self.device, dtype=mean_weight.dtype)
            l_h = factors["L_V"].to(device=self.device, dtype=mean_weight.dtype)

            weight_shape = factors["weight_shape"]
            z = torch.randn(
                weight_shape, device=self.device, dtype=mean_weight.dtype
            )

            sampled_weight = mean_weight + temperature * (l_h @ z @ l_q.T)

            prefix = f"{name}." if name else ""
            sampled_state[f"{prefix}weight"] = sampled_weight.detach().cpu()

            if module.bias is not None:
                sampled_state[f"{prefix}bias"] = module.bias.detach().cpu()

        model_sample = copy.deepcopy(self.model)
        model_sample.load_state_dict(sampled_state, strict=True)
        model_sample.to(self.device)
        model_sample.eval()
        samples.append(model_sample)

    return samples

def toggle_verbose(self) -> None:
    """Toggles verbosity flag."""
    self.verbose = not self.verbose
    print("Verbose:", "on" if self.verbose else "off")

def build_ensemble(self, n_members: int = 10) -> Ensemble:
    """Builds an Ensemble of networks sampled from the posterior.

    Args:
        n_members: Number of sampled members. Defaults to 10.

    Returns:
        Ensemble: Ensemble wrapping the sampled networks.
    """
    return Ensemble.from_posterior(self, n_members=n_members)

def fit(
    self,
    train_loader: DataLoader,
    val_loader: DataLoader | None = None,
    num_epochs: int = 100,
    step_clip: float | None = 2.0,
    prior_refresh: int = 1,
    **kwargs: Any,
) -> dict[str, list[float]]:
    """Runs assumed-density filtering over the training data.

    Each epoch visits every training point once in random order, updates
    the weight posteriors, then refreshes the noise and prior
    hyperparameters.

    Args:
        train_loader: DataLoader yielding (inputs, targets) pairs.
        val_loader: Optional DataLoader evaluated after every epoch.
        num_epochs: Number of passes over the training data. Defaults to 100.
        step_clip: Clipping threshold for each ADF update, or None to
            disable clipping. Defaults to 2.0.
        prior_refresh: Number of prior-refresh iterations per epoch, or 0
            to skip. Defaults to 1.
        **kwargs: Ignored, accepted for interface compatibility.

    Returns:
        dict[str, list[float]]: Per-epoch RMSE and NLPD on the training
        data, plus validation values when `val_loader` is given.
    """
    history: dict[str, list[float]] = {"train_rmse": [], "train_nlpd": []}
    if val_loader is not None:
        history["val_rmse"] = []
        history["val_nlpd"] = []

    for epoch in range(num_epochs):
        X_full, y_full = self._collect_dataset(train_loader)
        order = torch.randperm(X_full.shape[0], device=self.device)
        logZ_acc = torch.tensor(0.0, device=self.device, dtype=self.dtype)
        logZ1_acc = torch.tensor(0.0, device=self.device, dtype=self.dtype)
        logZ2_acc = torch.tensor(0.0, device=self.device, dtype=self.dtype)
        for idx in order.tolist():
            x = X_full[idx]
            y = y_full[idx]
            logZ, logZ1, logZ2 = self._single_datapoint_adf_step(x, y, step_clip)
            logZ_acc += logZ
            logZ1_acc += logZ1
            logZ2_acc += logZ2

        # Update noise hyperparameters once per epoch for stability.
        denom = max(len(order), 1)
        logZ_avg = logZ_acc / denom
        logZ1_avg = logZ1_acc / denom
        logZ2_avg = logZ2_acc / denom
        self.alpha_g, self.beta_g = self._gamma_adf_update_from_Z(
            logZ_avg, logZ1_avg, logZ2_avg, self.alpha_g, self.beta_g
        )

        if prior_refresh > 0:
            self._prior_refresh_epoch(n_refresh=prior_refresh, step_clip=step_clip)

        train_rmse, train_nlpd = self._evaluate_loader(train_loader)
        history["train_rmse"].append(train_rmse)
        history["train_nlpd"].append(train_nlpd)

        if val_loader is not None:
            val_rmse, val_nlpd = self._evaluate_loader(val_loader)
            history["val_rmse"].append(val_rmse)
            history["val_nlpd"].append(val_nlpd)

    self.is_fitted = True
    return history

def noise_variance(self) -> torch.Tensor:
    """Returns the posterior mean of the observation noise variance.

    Returns:
        torch.Tensor: Scalar noise variance under the current Gamma posterior.
    """
    alpha = torch.clamp(self.alpha_g, min=1.0 + 1e-6)
    return self.beta_g / (alpha - 1.0)

def sample_models(self, n_models: int = 10, **kwargs: Any) -> list[nn.Module]:
    """Draws deterministic networks from the fitted weight posterior.

    Args:
        n_models: Number of networks to sample. Defaults to 10.
        **kwargs: Ignored, accepted for interface compatibility.

    Returns:
        list[nn.Module]: Sampled networks in eval mode.

    Raises:
        RuntimeError: If `fit` has not been called.
    """
    if not self.is_fitted:
        raise RuntimeError("PBPEngine not fitted. Call fit() first.")

    models = []
    for _ in range(n_models):
        model_copy = self._sample_single_model()
        models.append(model_copy)
    return models

def forward_moments(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Propagate mean/variance through the network.

    Returns:
        (mz, vz): predictive mean and predictive variance of the network output,
        under the current factorized Gaussian approximation over weights.
    """
    x = x.to(device=self.device, dtype=self.dtype)
    assert x.dim() == 2 and x.shape[0] >= 1
    batch = x.shape[0]

    # Represent bias as an extra constant input dimension.
    mz = torch.cat(
        [x, torch.ones(batch, 1, device=self.device, dtype=self.dtype)], dim=1
    )  # (B, D+1)
    vz = torch.zeros_like(mz)

    for li, layer in enumerate(self.layers):
        d = layer.in_features + 1
        scale = 1.0 / math.sqrt(d)

        # Linear transform with weight scaling; ma/va are pre-activation moments.
        ma = (mz @ layer.m.t()) * scale  # (B, H)
        term1 = vz @ (layer.m**2).t()
        term2 = (mz**2) @ layer.v.t()
        term3 = vz @ layer.v.t()
        va = (term1 + term2 + term3) * (scale**2)

        is_last = li == len(self.layers) - 1
        if not is_last:
            # Nonlinearity is approximated by matching ReLU moments.
            mb, vb = relu_moments(ma, va)
            mz = torch.cat(
                [mb, torch.ones(batch, 1, device=self.device, dtype=self.dtype)],
                dim=1,
            )
            vz = torch.cat(
                [vb, torch.zeros(batch, 1, device=self.device, dtype=self.dtype)],
                dim=1,
            )
        else:
            mz = ma
            vz = va

    return mz, vz