VaLLA

laplace.baselaplace.VaLLA

VaLLA(
    model: Module,
    likelihood: Likelihood | str,
    inducing_locations: Tensor | MutableMapping | str,
    num_inducing: int | None = None,
    inducing_classes: Tensor | None = None,
    sigma_noise: float | Tensor = 1.0,
    prior_precision: float | Tensor = 1.0,
    prior_mean: float | Tensor = 0.0,
    temperature: float = 1.0,
    enable_backprop: bool = False,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    backend: type[CurvatureInterface] | None = None,
    backend_kwargs: dict[str, Any] | None = None,
    seed: int = 0,
    alpha: float = 1.0,
    mc_softmax_samples: int = 0,
)

Bases: BaseFunctionalLaplace

Variational linearized Laplace with inducing outputs and a fixed MAP model.

The network weights remain fixed while fit optimizes the variational factor, inducing locations, prior precision, and regression noise.

Parameters:

• model (Module) –

  Pretrained network to linearize.

• likelihood (Likelihood or str) –

  "classification", "regression", or "reward_modeling".

• inducing_locations (torch.Tensor, MutableMapping, or str) –

  Initial inputs, or "random" or "kmeans" to select them from the training loader. Floating-point tensor locations are optimized; mapping-style and integer tensor locations remain fixed.

• num_inducing (int or None, default: None ) –

  Number of locations required with a selection strategy. If explicit locations are supplied, this must match their count when provided.

• inducing_classes (Tensor or None, default: None ) –

  Output coordinate selected at each inducing location, shaped [num_inducing]. Explicit inputs default to coordinate zero. Strategies use training class labels for classification and coordinate zero for regression and individual rewards.

• sigma_noise (float or Tensor, default: 1.0 ) –

  Initial observation standard deviation for regression.

• prior_precision (float or Tensor, default: 1.0 ) –

  Initial scalar parameter-prior precision.

• prior_mean (float or Tensor, default: 0.0 ) –

  Prior mean; only zero is supported.

• temperature (float, default: 1.0 ) –

  Positive scale of the likelihood contribution.

• enable_backprop (bool, default: False ) –

  Retain derivatives through predictive moments.

• dict_key_x (str, default: "input_ids" ) –

  Input key for mapping-style batches.

• dict_key_y (str, default: "labels" ) –

  Target key for mapping-style batches.

• backend (type[CurvatureInterface] or None, default: None ) –

  Curvature backend used for parameter Jacobians.

• backend_kwargs (dict[str, Any] or None, default: None ) –

  Additional curvature backend options.

• seed (int, default: 0 ) –

  Seed for inducing-point selection and Monte Carlo fitting.

• alpha (float, default: 1.0 ) –

  Alpha-divergence value in [0, 1]; zero uses the ELBO limit.

• mc_softmax_samples (int, default: 0 ) –

  Monte Carlo draws for classification data terms when positive. Required for classification or reward modeling when alpha != 1.

Attributes:

• fit_history_ (dict[str, list[float]]) –

  objective records minibatch-scaled objective values computed before each update; val_nll records validation negative log likelihood.

• L (Parameter) –

  Packed lower-triangular variational factor containing num_inducing * (num_inducing + 1) // 2 entries.

• log_prior_precision (Parameter) –

  Trainable logarithm of scalar parameter-prior precision.

• log_noise_variance (Parameter or None) –

  Trainable logarithm of observation variance for regression; None for classification and reward modeling.

• generator (Generator) –

  Random generator for Monte Carlo data terms on the model device.

Notes

fit returns None and records objective and validation NLL in fit_history_. Repeated fits reset variational and inducing state by default; fit(..., override=False) continues optimization.

Source code in laplace/baselaplace.py

def __init__(
    self,
    model: torch.nn.Module,
    likelihood: Likelihood | str,
    inducing_locations: torch.Tensor | MutableMapping | str,
    num_inducing: int | None = None,
    inducing_classes: torch.Tensor | None = None,
    sigma_noise: float | torch.Tensor = 1.0,
    prior_precision: float | torch.Tensor = 1.0,
    prior_mean: float | torch.Tensor = 0.0,
    temperature: float = 1.0,
    enable_backprop: bool = False,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    backend: type[CurvatureInterface] | None = None,
    backend_kwargs: dict[str, Any] | None = None,
    seed: int = 0,
    alpha: float = 1.0,
    mc_softmax_samples: int = 0,
) -> None:
    """Initialize VaLLA from a pretrained model and inducing setup."""
    if torch.as_tensor(prior_precision).numel() != 1:
        raise ValueError("VaLLA requires scalar prior_precision.")
    if torch.any(torch.as_tensor(prior_mean) != 0):
        raise ValueError("VaLLA currently requires prior_mean=0.")
    if not isfinite(temperature) or temperature <= 0:
        raise ValueError("temperature must be positive and finite.")
    if not 0 <= alpha <= 1:
        raise ValueError("alpha must be in [0, 1].")
    if not isinstance(mc_softmax_samples, int) or mc_softmax_samples < 0:
        raise ValueError("mc_softmax_samples must be a nonnegative integer.")
    super().__init__(
        model,
        likelihood,
        sigma_noise=sigma_noise,
        prior_precision=prior_precision,
        prior_mean=prior_mean,
        temperature=temperature,
        enable_backprop=enable_backprop,
        dict_key_x=dict_key_x,
        dict_key_y=dict_key_y,
        backend=backend,
        backend_kwargs=backend_kwargs,
    )
    if (
        self.likelihood != Likelihood.REGRESSION
        and alpha != 1
        and mc_softmax_samples == 0
    ):
        raise ValueError(
            "Classification alpha != 1 requires mc_softmax_samples > 0; "
            "the deterministic probit approximation is alpha-independent."
        )
    self.seed = seed
    self.alpha = alpha
    self.mc_softmax_samples = mc_softmax_samples
    self.generator = torch.Generator(device=self._device).manual_seed(seed)
    self.log_prior_precision = torch.nn.Parameter(
        self._prior_precision.detach().log().clone()
    )
    self.log_noise_variance = (
        torch.nn.Parameter(2 * self._sigma_noise.detach().log().clone())
        if likelihood == Likelihood.REGRESSION
        else None
    )
    self._inducing_strategy = (
        inducing_locations if isinstance(inducing_locations, str) else None
    )
    if self._inducing_strategy is not None and self._inducing_strategy not in {
        "random",
        "kmeans",
    }:
        raise ValueError(
            "inducing_locations must be inputs, 'random', or 'kmeans'."
        )
    if isinstance(inducing_locations, str):
        if (
            not isinstance(num_inducing, Integral)
            or isinstance(num_inducing, bool)
            or num_inducing <= 0
        ):
            raise ValueError(
                "num_inducing must be a positive integer for an inducing strategy."
            )
        self.num_inducing = num_inducing
        self.inducing_locations = None
    else:
        self.inducing_locations = self._make_inducing(inducing_locations)
        count = (
            self.inducing_locations[self.dict_key_x].shape[0]
            if isinstance(self.inducing_locations, MutableMapping)
            else self.inducing_locations.shape[0]
        )
        if num_inducing is not None and (
            not isinstance(num_inducing, Integral)
            or isinstance(num_inducing, bool)
            or num_inducing != count
        ):
            raise ValueError("num_inducing must match inducing_locations.")
        self.num_inducing = count
    if self.num_inducing <= 0:
        raise ValueError("At least one inducing location is required.")
    self.inducing_classes = (
        torch.zeros(self.num_inducing, device=self._device, dtype=torch.long)
        if inducing_classes is None
        else torch.as_tensor(
            inducing_classes, device=self._device, dtype=torch.long
        )
    )
    if self.inducing_classes.shape != (self.num_inducing,):
        raise ValueError("inducing_classes must have one entry per inducing input.")
    self._initial_inducing_locations = (
        None
        if self._inducing_strategy is not None
        else self._clone_fixed_inputs(self.inducing_locations)
    )
    self._initial_inducing_classes = self.inducing_classes.detach().clone()
    self._init_variational_factor()
    self.n_data = 0
    self._fitted = False
    self.fit_history_: dict[str, list[float]] = {"objective": [], "val_nll": []}

model

model: Module = model

likelihood

likelihood: Likelihood | str = likelihood

params

params: list[Tensor] = []

is_subset_params

is_subset_params: bool = False

n_params

n_params: int = sum(numel() for p in params)

n_layers

n_layers: int = len(params)

enable_backprop

enable_backprop: bool = enable_backprop

dict_key_x

dict_key_x = dict_key_x

dict_key_y

dict_key_y = dict_key_y

_backend

_backend: CurvatureInterface | None = None

_backend_cls

_backend_cls: type[CurvatureInterface] = backend

_backend_kwargs

_backend_kwargs: dict[str, Any] = (
    dict() if backend_kwargs is None else backend_kwargs
)

_asdl_fisher_kwargs

_asdl_fisher_kwargs: dict[str, Any] = (
    dict()
    if asdl_fisher_kwargs is None
    else asdl_fisher_kwargs
)

loss

loss: float = 0.0

n_outputs

n_outputs: int = 0

_prior_mean

_prior_mean: Tensor

_prior_precision

_prior_precision: Tensor

_sigma_noise

_sigma_noise: Tensor

_posterior_scale

_posterior_scale: Tensor | None

_device

_device: device

_dtype

_dtype: dtype

backend

backend: CurvatureInterface

prior_precision_diag

prior_precision_diag: Tensor

Obtain the diagonal prior precision \(p_0\) constructed from either a scalar, layer-wise, or diagonal prior precision.

Returns:

• prior_precision_diag ( Tensor ) –

_H_factor

_H_factor: Tensor

_uses_backpack_backend

_uses_backpack_backend: bool

Whether the curvature backend extends modules using BackPACK.

_uses_asdl_backend

_uses_asdl_backend: bool

Whether the curvature backend computes Jacobians using ASDL.

seed

seed = seed

alpha

alpha = alpha

mc_softmax_samples

mc_softmax_samples = mc_softmax_samples

generator

generator = manual_seed(seed)

log_prior_precision

log_prior_precision = Parameter(clone())

log_noise_variance

log_noise_variance = (
    Parameter(2 * clone())
    if likelihood == REGRESSION
    else None
)

_inducing_strategy

_inducing_strategy = (
    inducing_locations
    if isinstance(inducing_locations, str)
    else None
)

num_inducing

num_inducing = num_inducing

inducing_locations

inducing_locations = None

inducing_classes

inducing_classes = (
    zeros(num_inducing, device=_device, dtype=long)
    if inducing_classes is None
    else as_tensor(
        inducing_classes, device=_device, dtype=long
    )
)

_initial_inducing_locations

_initial_inducing_locations = (
    None
    if _inducing_strategy is not None
    else _clone_fixed_inputs(inducing_locations)
)

_initial_inducing_classes

_initial_inducing_classes = clone()

n_data

n_data = 0

_fitted

_fitted = False

fit_history_

fit_history_: dict[str, list[float]] = {
    "objective": [],
    "val_nll": [],
}

prior_precision

prior_precision: Tensor

Return the positive variational prior precision.

prior_mean

prior_mean: Tensor

Return the zero parameter-prior mean.

temperature

temperature: float

Return the positive likelihood temperature.

sigma_noise

sigma_noise: Tensor

Return the regression observation standard deviation.

log_likelihood

log_likelihood: Tensor

Reject access to an untracked training log likelihood.

Raises:

• NotImplementedError –

  Always raised because VaLLA records its alpha-divergence objective rather than a training log likelihood.

_curv_closure

_curv_closure(
    X: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    N: int,
) -> tuple[Tensor, Tensor]

Source code in laplace/baselaplace.py

def _curv_closure(
    self,
    X: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    N: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    raise NotImplementedError

__call__

__call__(
    x: Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GP,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: Generator | None = None,
    fitting: bool = False,
    **model_kwargs: dict[str, Any]
) -> Tensor | tuple[Tensor, Tensor]

Predict from the fitted function-space posterior.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• pred_type (PredType or str, default: PredType.GP ) –

  Prediction type. Function-space estimators support only "gp".

• joint (bool, default: False ) –

  Retain cross-input covariance for regression predictions.

• link_approx (LinkApprox or str, default: LinkApprox.PROBIT ) –

  Approximation for the classification probability link.

• n_samples (int, default: 100 ) –

  Number of draws for a Monte Carlo classification link.

• diagonal_output (bool, default: False ) –

  Return only output variances for regression when supported.

• generator (Generator or None, default: None ) –

  Random generator for Monte Carlo classification predictions.

• fitting (bool, default: False ) –

  Evaluate reward-model preference pairs as classification inputs.

• **model_kwargs (Any, default: {} ) –

  Accepted for compatibility with the base prediction interface; these arguments are not forwarded to the model.

Returns:

• Tensor or tuple[Tensor, Tensor] –

  Classification probabilities [B, C], or regression/reward latent mean and covariance. With joint=True, the regression mean is flattened and covariance has shape [B*O, B*O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty, the prediction type or link is unsupported, or the generator has an incompatible type or device.

Notes

Regression and individual reward predictions exclude observation noise. Their marginal mean is [B, O] and covariance is [B, O, O], or variance [B, O] with diagonal_output=True.

Source code in laplace/baselaplace.py

def __call__(
    self,
    x: torch.Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GP,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
    fitting: bool = False,
    **model_kwargs: dict[str, Any],
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Predict from the fitted function-space posterior.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    pred_type : PredType or str, default=PredType.GP
        Prediction type. Function-space estimators support only ``"gp"``.
    joint : bool, default=False
        Retain cross-input covariance for regression predictions.
    link_approx : LinkApprox or str, default=LinkApprox.PROBIT
        Approximation for the classification probability link.
    n_samples : int, default=100
        Number of draws for a Monte Carlo classification link.
    diagonal_output : bool, default=False
        Return only output variances for regression when supported.
    generator : torch.Generator or None, default=None
        Random generator for Monte Carlo classification predictions.
    fitting : bool, default=False
        Evaluate reward-model preference pairs as classification inputs.
    **model_kwargs : Any
        Accepted for compatibility with the base prediction interface;
        these arguments are not forwarded to the model.

    Returns
    -------
    torch.Tensor or tuple[torch.Tensor, torch.Tensor]
        Classification probabilities ``[B, C]``, or regression/reward
        latent mean and covariance. With ``joint=True``, the regression
        mean is flattened and covariance has shape ``[B*O, B*O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty, the prediction type or link is
        unsupported, or the generator has an incompatible type or device.

    Notes
    -----
    Regression and individual reward predictions exclude observation
    noise. Their marginal mean is ``[B, O]`` and covariance is
    ``[B, O, O]``, or variance ``[B, O]`` with ``diagonal_output=True``.
    """
    self._check_fitted()
    self._check_nonempty_inputs(x)
    if pred_type != PredType.GP:
        raise ValueError("Only gp is supported as a functional prediction type.")
    if link_approx not in list(LinkApprox):
        raise ValueError(f"Unsupported link approximation {link_approx}.")
    if generator is not None and (
        not isinstance(generator, torch.Generator)
        or generator.device != self._device
    ):
        raise ValueError("Invalid random generator (check type and device).")
    likelihood = self.likelihood
    if likelihood == Likelihood.REWARD_MODELING:
        likelihood = Likelihood.CLASSIFICATION if fitting else Likelihood.REGRESSION
    if likelihood == Likelihood.CLASSIFICATION and link_approx == LinkApprox.MC:
        return self.predictive_samples(
            x,
            n_samples=n_samples,
            diagonal_output=diagonal_output,
            generator=generator,
            fitting=fitting,
        ).mean(dim=0)
    return self._glm_forward_call(
        x, likelihood, joint, link_approx, n_samples, diagonal_output
    )

predictive

predictive(
    x: Tensor,
    pred_type: PredType | str,
    link_approx: LinkApprox | str,
    n_samples: int,
) -> Tensor | tuple[Tensor, Tensor]

Source code in laplace/baselaplace.py

def predictive(
    self,
    x: torch.Tensor,
    pred_type: PredType | str,
    link_approx: LinkApprox | str,
    n_samples: int,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    return self(
        x, pred_type=pred_type, link_approx=link_approx, n_samples=n_samples
    )

_check_jacobians

_check_jacobians(Js: Tensor) -> None

Source code in laplace/baselaplace.py

def _check_jacobians(self, Js: torch.Tensor) -> None:
    if not isinstance(Js, torch.Tensor):
        raise ValueError("Jacobians have to be torch.Tensor.")
    if not Js.device == self._device:
        raise ValueError("Jacobians need to be on the same device as Laplace.")
    m, k, p = Js.size()
    if p != self.n_params:
        raise ValueError("Invalid Jacobians shape for Laplace posterior approx.")

_gridsearch

_gridsearch(
    loss: Metric | Callable[[Tensor], Tensor | float],
    interval: Tensor,
    val_loader: DataLoader,
    pred_type: PredType | str,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    progress_bar: bool = False,
) -> Tensor

Source code in laplace/baselaplace.py

def _gridsearch(
    self,
    loss: torchmetrics.Metric | Callable[[torch.Tensor], torch.Tensor | float],
    interval: torch.Tensor,
    val_loader: DataLoader,
    pred_type: PredType | str,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    progress_bar: bool = False,
) -> torch.Tensor:
    assert callable(loss) or isinstance(loss, torchmetrics.Metric)

    results: list[float] = list()
    prior_precs: list[torch.Tensor] = list()
    pbar = tqdm.tqdm(interval, disable=not progress_bar)

    for prior_prec in pbar:
        self.prior_precision = prior_prec

        try:
            result = validate(
                self,
                val_loader,
                loss,
                pred_type=pred_type,
                link_approx=link_approx,
                n_samples=n_samples,
                dict_key_y=self.dict_key_y,
            )
        except LinAlgError:
            result = np.inf
        except RuntimeError as err:
            if "not positive definite" in str(err):
                result = np.inf
            else:
                raise err

        if progress_bar:
            pbar.set_description(
                f"[Grid search | prior_prec: {prior_prec:.3e}, loss: {result:.3f}]"
            )

        results.append(result)
        prior_precs.append(prior_prec)

    return prior_precs[np.argmin(results)]

_glm_forward_call

_glm_forward_call(
    x: Tensor | MutableMapping,
    likelihood: Likelihood | str,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
) -> Tensor | tuple[Tensor, Tensor]

Compute the posterior predictive on input data x for "glm" pred type.

Parameters:

• x (Tensor or MutableMapping) –

  (batch_size, input_shape) if tensor. If MutableMapping, must contain the said tensor.

• likelihood (Likelihood or str in {'classification', 'regression', 'reward_modeling'}) –

  determines the log likelihood Hessian approximation.

• link_approx (('mc', 'probit', 'bridge', 'bridge_norm'), default: 'mc' ) –

  how to approximate the classification link function for the 'glm'. For pred_type='nn', only 'mc' is possible.

• joint (bool, default: False ) –

  Whether to output a joint predictive distribution in regression with pred_type='glm'. If set to True, the predictive distribution has the same form as GP posterior, i.e. N([f(x1), ...,f(xm)], Cov[f(x1), ..., f(xm)]). If False, then only outputs the marginal predictive distribution. Only available for regression and GLM predictive.

• n_samples (int, default: 100 ) –

  number of samples for link_approx='mc'.

• diagonal_output (bool, default: False ) –

  whether to use a diagonalized posterior predictive on the outputs. Only works for pred_type='glm' and link_approx='mc'.

Returns:

• predictive ( Tensor or tuple[Tensor] ) –

  For likelihood='classification', a torch.Tensor is returned with a distribution over classes (similar to a Softmax). For likelihood='regression', a tuple of torch.Tensor is returned with the mean and the predictive variance. For likelihood='regression' and joint=True, a tuple of torch.Tensor is returned with the mean and the predictive covariance.

Source code in laplace/baselaplace.py

def _glm_forward_call(
    self,
    x: torch.Tensor | MutableMapping,
    likelihood: Likelihood | str,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Compute the posterior predictive on input data `x` for "glm" pred type.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        `(batch_size, input_shape)` if tensor. If MutableMapping, must contain
        the said tensor.

    likelihood : Likelihood or str in {'classification', 'regression', 'reward_modeling'}
        determines the log likelihood Hessian approximation.

    link_approx : {'mc', 'probit', 'bridge', 'bridge_norm'}
        how to approximate the classification link function for the `'glm'`.
        For `pred_type='nn'`, only 'mc' is possible.

    joint : bool
        Whether to output a joint predictive distribution in regression with
        `pred_type='glm'`. If set to `True`, the predictive distribution
        has the same form as GP posterior, i.e. N([f(x1), ...,f(xm)], Cov[f(x1), ..., f(xm)]).
        If `False`, then only outputs the marginal predictive distribution.
        Only available for regression and GLM predictive.

    n_samples : int
        number of samples for `link_approx='mc'`.

    diagonal_output : bool
        whether to use a diagonalized posterior predictive on the outputs.
        Only works for `pred_type='glm'` and `link_approx='mc'`.

    Returns
    -------
    predictive: torch.Tensor or tuple[torch.Tensor]
        For `likelihood='classification'`, a torch.Tensor is returned with
        a distribution over classes (similar to a Softmax).
        For `likelihood='regression'`, a tuple of torch.Tensor is returned
        with the mean and the predictive variance.
        For `likelihood='regression'` and `joint=True`, a tuple of torch.Tensor
        is returned with the mean and the predictive covariance.
    """
    f_mu, f_var = self._glm_predictive_distribution(
        x, joint=joint and likelihood == Likelihood.REGRESSION
    )

    if likelihood == Likelihood.REGRESSION:
        if diagonal_output and not joint:
            f_var = torch.diagonal(f_var, dim1=-2, dim2=-1)
        return f_mu, f_var

    if link_approx == LinkApprox.MC:
        return self._glm_predictive_samples(
            f_mu,
            f_var,
            n_samples=n_samples,
            diagonal_output=diagonal_output,
        ).mean(dim=0)
    elif link_approx == LinkApprox.PROBIT:
        kappa = 1 / torch.sqrt(1.0 + np.pi / 8 * f_var.diagonal(dim1=1, dim2=2))
        return torch.softmax(kappa * f_mu, dim=-1)
    elif "bridge" in link_approx:
        # zero mean correction
        f_mu -= (
            f_var.sum(-1)
            * f_mu.sum(-1).reshape(-1, 1)
            / f_var.sum(dim=(1, 2)).reshape(-1, 1)
        )
        f_var -= torch.einsum(
            "bi,bj->bij", f_var.sum(-1), f_var.sum(-2)
        ) / f_var.sum(dim=(1, 2)).reshape(-1, 1, 1)

        # Laplace Bridge
        _, K = f_mu.size(0), f_mu.size(-1)
        f_var_diag = torch.diagonal(f_var, dim1=1, dim2=2)

        # optional: variance correction
        if link_approx == LinkApprox.BRIDGE_NORM:
            f_var_diag_mean = f_var_diag.mean(dim=1)
            f_var_diag_mean /= torch.as_tensor(
                [K / 2], device=self._device, dtype=self._dtype
            ).sqrt()
            f_mu /= f_var_diag_mean.sqrt().unsqueeze(-1)
            f_var_diag /= f_var_diag_mean.unsqueeze(-1)

        sum_exp = torch.exp(-f_mu).sum(dim=1).unsqueeze(-1)
        alpha = (1 - 2 / K + f_mu.exp() / K**2 * sum_exp) / f_var_diag
        return torch.nan_to_num(alpha / alpha.sum(dim=1).unsqueeze(-1), nan=1.0)
    else:
        raise ValueError(
            "Prediction path invalid. Check the likelihood, pred_type, link_approx combination!"
        )

sample

sample(
    n_samples: int = 1, generator: Generator | None = None
) -> Tensor

Sample from the Laplace posterior approximation, i.e., \( \theta \sim \mathcal{N}(\theta_{MAP}, P^{-1})\).

Parameters:

• n_samples (int, default: 100 ) –

  number of samples

• generator (Generator, default: None ) –

  random number generator to control the samples

Returns:

• samples ( Tensor ) –

Source code in laplace/baselaplace.py

def sample(
    self, n_samples: int = 1, generator: torch.Generator | None = None
) -> torch.Tensor:
    """Sample from the Laplace posterior approximation, i.e.,
    \\( \\theta \\sim \\mathcal{N}(\\theta_{MAP}, P^{-1})\\).

    Parameters
    ----------
    n_samples : int, default=100
        number of samples

    generator : torch.Generator, optional
        random number generator to control the samples

    Returns
    -------
    samples: torch.Tensor
    """
    raise NotImplementedError

functional_samples

functional_samples(
    x: Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: Generator | None = None,
    joint: bool = False,
) -> Tensor

Draw latent outputs; use joint=True to retain cross-input covariance.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• pred_type (PredType or str, default: PredType.GLM ) –

  Accepted for API compatibility. Every value draws from the function-space posterior; "nn" does not draw model weights.

• n_samples (int, default: 100 ) –

  Number of draws.

• diagonal_output (bool, default: False ) –

  Drop cross-output covariance before sampling.

• generator (Generator or None, default: None ) –

  Random generator on the model device.

• joint (bool, default: False ) –

  Retain covariance between input rows.

Returns:

• Tensor –

  Latent draws shaped [n_samples, B, O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty, the sampling type is unsupported, or n_samples is not positive.

Source code in laplace/baselaplace.py

def functional_samples(
    self,
    x: torch.Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
    joint: bool = False,
) -> torch.Tensor:
    """Draw latent outputs; use `joint=True` to retain cross-input covariance.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    pred_type : PredType or str, default=PredType.GLM
        Accepted for API compatibility. Every value draws from the
        function-space posterior; ``"nn"`` does not draw model weights.
    n_samples : int, default=100
        Number of draws.
    diagonal_output : bool, default=False
        Drop cross-output covariance before sampling.
    generator : torch.Generator or None, default=None
        Random generator on the model device.
    joint : bool, default=False
        Retain covariance between input rows.

    Returns
    -------
    torch.Tensor
        Latent draws shaped ``[n_samples, B, O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty, the sampling type is unsupported,
        or ``n_samples`` is not positive.
    """
    self._check_fitted()
    self._check_nonempty_inputs(x)
    if pred_type not in list(PredType):
        raise ValueError(f"Unsupported functional sampling type {pred_type}.")
    if n_samples <= 0:
        raise ValueError("n_samples must be positive.")
    mean, covariance = self._glm_predictive_distribution(x, joint=joint)
    if joint:
        batch = (
            x[self.dict_key_x].shape[0]
            if isinstance(x, MutableMapping)
            else x.shape[0]
        )
        n_outputs = mean.numel() // batch
        if diagonal_output:
            output_ids = torch.arange(mean.numel(), device=mean.device) % n_outputs
            covariance = covariance * (output_ids[:, None] == output_ids[None, :])
        draws = draw_gaussian(mean, covariance, n_samples, generator)
        return draws.reshape(n_samples, batch, n_outputs)
    return self._glm_functional_samples(
        mean, covariance, n_samples, diagonal_output, generator
    )

_glm_functional_samples

_glm_functional_samples(
    f_mu: Tensor,
    f_var: Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: Generator | None = None,
) -> Tensor

Draw per-input latent samples from supplied Gaussian moments.

Parameters:

• f_mu (Tensor) –

  Latent means shaped [B, O].

• f_var (Tensor) –

  Per-input covariance shaped [B, O, O].

• n_samples (int) –

  Number of draws.

• diagonal_output (bool, default: False ) –

  Drop covariance between output coordinates before sampling.

• generator (Generator or None, default: None ) –

  Random generator for the draws.

Returns:

• Tensor –

  Latent draws shaped [n_samples, B, O].

Source code in laplace/baselaplace.py

def _glm_functional_samples(
    self,
    f_mu: torch.Tensor,
    f_var: torch.Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    """Draw per-input latent samples from supplied Gaussian moments.

    Parameters
    ----------
    f_mu : torch.Tensor
        Latent means shaped ``[B, O]``.
    f_var : torch.Tensor
        Per-input covariance shaped ``[B, O, O]``.
    n_samples : int
        Number of draws.
    diagonal_output : bool, default=False
        Drop covariance between output coordinates before sampling.
    generator : torch.Generator or None, default=None
        Random generator for the draws.

    Returns
    -------
    torch.Tensor
        Latent draws shaped ``[n_samples, B, O]``.
    """
    covariance = (
        torch.diag_embed(f_var.diagonal(dim1=-2, dim2=-1))
        if diagonal_output
        else f_var
    )
    return draw_gaussian(f_mu, covariance, n_samples, generator)

_glm_predictive_samples

_glm_predictive_samples(
    f_mu: Tensor,
    f_var: Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: Generator | None = None,
) -> Tensor

Sample from the posterior predictive on input data x using "glm" prediction type. I.e., the inverse-link function correponding to the likelihood is applied on top of the functional sample.

Parameters:

• f_mu (Tensor or MutableMapping) –

  glm predictive mean (batch_size, output_shape)

• f_var (Tensor or MutableMapping) –

  glm predictive covariances (batch_size, output_shape, output_shape)

• n_samples (int) –

  number of samples

• diagonal_output (bool, default: False ) –

  whether to use a diagonalized glm posterior predictive on the outputs.

• generator (Generator, default: None ) –

  random number generator to control the samples (if sampling used)

Returns:

• samples ( Tensor ) –

  samples (n_samples, batch_size, output_shape)

Source code in laplace/baselaplace.py

def _glm_predictive_samples(
    self,
    f_mu: torch.Tensor,
    f_var: torch.Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    """Sample from the posterior predictive on input data `x` using "glm" prediction
    type. I.e., the inverse-link function correponding to the likelihood is applied
    on top of the functional sample.

    Parameters
    ----------
    f_mu : torch.Tensor or MutableMapping
        glm predictive mean `(batch_size, output_shape)`

    f_var : torch.Tensor or MutableMapping
        glm predictive covariances `(batch_size, output_shape, output_shape)`

    n_samples : int
        number of samples

    diagonal_output : bool
        whether to use a diagonalized glm posterior predictive on the outputs.

    generator : torch.Generator, optional
        random number generator to control the samples (if sampling used)

    Returns
    -------
    samples : torch.Tensor
        samples `(n_samples, batch_size, output_shape)`
    """
    f_samples = self._glm_functional_samples(
        f_mu, f_var, n_samples, diagonal_output, generator
    )

    if self.likelihood == Likelihood.REGRESSION:
        return f_samples
    else:
        return torch.softmax(f_samples, dim=-1)

_checkpoint_configuration

_checkpoint_configuration() -> dict[str, Any]

Record input keys and backend options required by a checkpoint.

Returns:

• dict[str, Any] –

  Settings that must match when the checkpoint is restored.

Source code in laplace/baselaplace.py

def _checkpoint_configuration(self) -> dict[str, Any]:
    """Record input keys and backend options required by a checkpoint.

    Returns
    -------
    dict[str, Any]
        Settings that must match when the checkpoint is restored.
    """
    return {
        "dict_key_x": self.dict_key_x,
        "dict_key_y": self.dict_key_y,
        "backend_class": (
            f"{self._backend_cls.__module__}.{self._backend_cls.__qualname__}"
        ),
        "backend_kwargs": deepcopy(self._backend_kwargs),
        "enable_backprop": self.enable_backprop,
    }

_validate_checkpoint_configuration

_validate_checkpoint_configuration(
    state_dict: dict[str, Any]
) -> None

Reject a checkpoint created with different input or backend settings.

Parameters:

• state_dict (dict[str, Any]) –

  Saved functional posterior state.

Raises:

• ValueError –

  If a required setting differs from the current estimator.

Source code in laplace/baselaplace.py

def _validate_checkpoint_configuration(self, state_dict: dict[str, Any]) -> None:
    """Reject a checkpoint created with different input or backend settings.

    Parameters
    ----------
    state_dict : dict[str, Any]
        Saved functional posterior state.

    Raises
    ------
    ValueError
        If a required setting differs from the current estimator.
    """
    for key, current in self._checkpoint_configuration().items():
        if not same_configuration(state_dict.get(key), current):
            raise ValueError(f"Checkpoint requires matching {key}.")

_refresh_posterior

_refresh_posterior() -> None

Update cached posterior quantities after a hyperparameter change.

Source code in laplace/baselaplace.py

def _refresh_posterior(self) -> None:
    """Update cached posterior quantities after a hyperparameter change."""

_check_fitted

_check_fitted() -> None

Refresh a fitted posterior or reject predictions before fitting.

Raises:

• RuntimeError –

  If fit has not completed.

Source code in laplace/baselaplace.py

def _check_fitted(self) -> None:
    """Refresh a fitted posterior or reject predictions before fitting.

    Raises
    ------
    RuntimeError
        If ``fit`` has not completed.
    """
    if not getattr(self, "_fitted", False):
        raise RuntimeError("Call fit before making functional Laplace predictions.")
    self._refresh_posterior()

_check_nonempty_inputs

_check_nonempty_inputs(x: Tensor | MutableMapping) -> None

Require at least one input row for a function-space prediction.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

Raises:

• ValueError –

  If the input batch is empty.

Source code in laplace/baselaplace.py

def _check_nonempty_inputs(self, x: torch.Tensor | MutableMapping) -> None:
    """Require at least one input row for a function-space prediction.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.

    Raises
    ------
    ValueError
        If the input batch is empty.
    """
    inputs = x[self.dict_key_x] if isinstance(x, MutableMapping) else x
    if inputs.shape[0] == 0:
        raise ValueError("Function-space predictions require at least one input.")

predictive_moments

predictive_moments(
    x: Tensor | MutableMapping, *, joint: bool = False
) -> tuple[Tensor, Tensor]

Return latent predictive moments without observation noise.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• joint (bool, default: False ) –

  Include covariance between input rows when True.

Returns:

• tuple[Tensor, Tensor] –

  Mean [B, O] and covariance [B, O, O] by default. With joint=True, the mean is flattened and covariance has shape [B*O, B*O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty.

Source code in laplace/baselaplace.py

def predictive_moments(
    self, x: torch.Tensor | MutableMapping, *, joint: bool = False
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return latent predictive moments without observation noise.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    joint : bool, default=False
        Include covariance between input rows when ``True``.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        Mean ``[B, O]`` and covariance ``[B, O, O]`` by default. With
        ``joint=True``, the mean is flattened and covariance has shape
        ``[B*O, B*O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty.
    """
    self._check_fitted()
    self._check_nonempty_inputs(x)
    return self._glm_predictive_distribution(x, joint=joint)

_classification_validation_nll

_classification_validation_nll(
    x: Tensor | MutableMapping, targets: Tensor
) -> Tensor

Compute summed classification NLL using the probit logit scale.

Parameters:

• x (Tensor or MutableMapping) –

  Validation inputs.

• targets (Tensor) –

  Class targets for those inputs.

Returns:

• Tensor –

  Summed negative log likelihood for the batch.

Source code in laplace/baselaplace.py

def _classification_validation_nll(
    self, x: torch.Tensor | MutableMapping, targets: torch.Tensor
) -> torch.Tensor:
    """Compute summed classification NLL using the probit logit scale.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Validation inputs.
    targets : torch.Tensor
        Class targets for those inputs.

    Returns
    -------
    torch.Tensor
        Summed negative log likelihood for the batch.
    """
    mean, covariance = self.predictive_moments(x)
    scale = torch.sqrt(1 + torch.pi / 8 * covariance.diagonal(dim1=-2, dim2=-1))
    return torch.nn.functional.cross_entropy(
        mean / scale, classification_targets(targets), reduction="sum"
    )

predictive_samples

predictive_samples(
    x: Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: Generator | None = None,
    fitting: bool = False,
    joint: bool = False,
) -> Tensor

Draw outputs and apply the classification link when appropriate.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• pred_type (PredType or str, default: PredType.GLM ) –

  Accepted for function-space sampling API compatibility.

• n_samples (int, default: 100 ) –

  Number of draws.

• diagonal_output (bool, default: False ) –

  Drop cross-output covariance before sampling.

• generator (Generator or None, default: None ) –

  Random generator on the model device.

• fitting (bool, default: False ) –

  Interpret reward-model inputs as preference pairs.

• joint (bool, default: False ) –

  Retain covariance between input rows.

Returns:

• Tensor –

  Predictive draws shaped [n_samples, B, O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty, the sampling type is unsupported, or n_samples is not positive.

Notes

Classification draws are class probabilities obtained by applying softmax to latent samples. Regression and individual reward draws are latent function values without observation noise.

Source code in laplace/baselaplace.py

def predictive_samples(
    self,
    x: torch.Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
    fitting: bool = False,
    joint: bool = False,
) -> torch.Tensor:
    """Draw outputs and apply the classification link when appropriate.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    pred_type : PredType or str, default=PredType.GLM
        Accepted for function-space sampling API compatibility.
    n_samples : int, default=100
        Number of draws.
    diagonal_output : bool, default=False
        Drop cross-output covariance before sampling.
    generator : torch.Generator or None, default=None
        Random generator on the model device.
    fitting : bool, default=False
        Interpret reward-model inputs as preference pairs.
    joint : bool, default=False
        Retain covariance between input rows.

    Returns
    -------
    torch.Tensor
        Predictive draws shaped ``[n_samples, B, O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty, the sampling type is unsupported,
        or ``n_samples`` is not positive.

    Notes
    -----
    Classification draws are class probabilities obtained by applying
    softmax to latent samples. Regression and individual reward draws are
    latent function values without observation noise.
    """
    samples = self.functional_samples(
        x, pred_type, n_samples, diagonal_output, generator, joint
    )
    if self.likelihood == Likelihood.CLASSIFICATION or (
        self.likelihood == Likelihood.REWARD_MODELING and fitting
    ):
        return samples.softmax(dim=-1)
    return samples

_clone_fixed_inputs

_clone_fixed_inputs(value: Any) -> Any

Detach and recursively copy fixed inducing inputs.

Parameters:

• value (Any) –

  Tensor, mapping, sequence, or metadata to copy.

Returns:

• Any –

  Independent copy with detached tensors and the same nesting.

Source code in laplace/baselaplace.py

@staticmethod
def _clone_fixed_inputs(value: Any) -> Any:
    """Detach and recursively copy fixed inducing inputs.

    Parameters
    ----------
    value : Any
        Tensor, mapping, sequence, or metadata to copy.

    Returns
    -------
    Any
        Independent copy with detached tensors and the same nesting.
    """
    if isinstance(value, torch.Tensor):
        return value.detach().clone()
    if isinstance(value, MutableMapping):
        return {key: VaLLA._clone_fixed_inputs(item) for key, item in value.items()}
    if isinstance(value, (tuple, list)):
        return type(value)(VaLLA._clone_fixed_inputs(item) for item in value)
    return deepcopy(value)

_make_inducing

_make_inducing(inputs: Any) -> Tensor | dict

Prepare inducing inputs and expose eligible locations for optimization.

Parameters:

• inputs (Any) –

  Batched inducing tensor, mapping-style model input, or an array-like input convertible to a floating-point tensor.

Returns:

• Tensor or dict –

  Copied inputs on the model device. Floating-point tensor inputs become trainable parameters in the model dtype. Mappings and integer tensors remain detached and fixed.

Source code in laplace/baselaplace.py

def _make_inducing(self, inputs: Any) -> torch.Tensor | dict:
    """Prepare inducing inputs and expose eligible locations for optimization.

    Parameters
    ----------
    inputs : Any
        Batched inducing tensor, mapping-style model input, or an
        array-like input convertible to a floating-point tensor.

    Returns
    -------
    torch.Tensor or dict
        Copied inputs on the model device. Floating-point tensor inputs
        become trainable parameters in the model dtype. Mappings and
        integer tensors remain detached and fixed.
    """
    converted = to_device(inputs, self._device, self._dtype)
    if isinstance(converted, MutableMapping):
        return self._clone_fixed_inputs(converted)
    if not isinstance(converted, torch.Tensor):
        converted = torch.as_tensor(
            converted, device=self._device, dtype=self._dtype
        )
    if converted.is_floating_point():
        return torch.nn.Parameter(converted.detach().clone())
    return converted.detach().clone()

_init_variational_factor

_init_variational_factor() -> None

Initialize the packed lower-triangular variational factor.

Source code in laplace/baselaplace.py

def _init_variational_factor(self) -> None:
    """Initialize the packed lower-triangular variational factor."""
    rows, cols = torch.tril_indices(self.num_inducing, self.num_inducing)
    initial = torch.eye(self.num_inducing, device=self._device, dtype=self._dtype)
    self.L = torch.nn.Parameter(initial[rows, cols].clone())

_factor_matrix

_factor_matrix() -> Tensor

Unpack the trainable lower-triangular variational factor.

Returns:

• Tensor –

  Lower-triangular matrix shaped [num_inducing, num_inducing] with gradients retained through its packed entries in L.

Source code in laplace/baselaplace.py

def _factor_matrix(self) -> torch.Tensor:
    """Unpack the trainable lower-triangular variational factor.

    Returns
    -------
    torch.Tensor
        Lower-triangular matrix shaped ``[num_inducing, num_inducing]``
        with gradients retained through its packed entries in ``L``.
    """
    rows, cols = torch.tril_indices(
        self.num_inducing, self.num_inducing, device=self._device
    )
    matrix = torch.zeros(
        self.num_inducing, self.num_inducing, device=self._device, dtype=self._dtype
    )
    matrix[rows, cols] = self.L
    return matrix

_batch

_batch(batch: Any) -> tuple[Any, Tensor]

Move a training batch and its targets to the model device.

Parameters:

• batch (Any) –

  (inputs, targets) or a mapping containing dict_key_y.

Returns:

• tuple[Any, Tensor] –

  Model inputs and targets on the model device. Floating-point inputs use the model dtype; target dtypes are preserved.

Source code in laplace/baselaplace.py

def _batch(self, batch: Any) -> tuple[Any, torch.Tensor]:
    """Move a training batch and its targets to the model device.

    Parameters
    ----------
    batch : Any
        ``(inputs, targets)`` or a mapping containing ``dict_key_y``.

    Returns
    -------
    tuple[Any, torch.Tensor]
        Model inputs and targets on the model device. Floating-point
        inputs use the model dtype; target dtypes are preserved.
    """
    x, y = split_batch(batch, self.dict_key_y)
    return to_device(x, self._device, self._dtype), y.to(self._device)

_gather_inducing

_gather_inducing(loader: DataLoader) -> None

Choose inducing inputs and output coordinates from training data.

Parameters:

• loader (DataLoader) –

  Training loader used for random or k-means selection.

Raises:

• ValueError –

  If the loader is empty, too small, or incompatible with k-means.

Source code in laplace/baselaplace.py

def _gather_inducing(self, loader: DataLoader) -> None:
    """Choose inducing inputs and output coordinates from training data.

    Parameters
    ----------
    loader : DataLoader
        Training loader used for random or k-means selection.

    Raises
    ------
    ValueError
        If the loader is empty, too small, or incompatible with k-means.
    """
    candidates = []
    classes = []
    source = loader
    if self._inducing_strategy == "random":
        available_indices = training_indices(loader)
        n_pairs = (self.num_inducing + 1) // 2
        needed = (
            n_pairs
            if self.likelihood == Likelihood.REWARD_MODELING
            else self.num_inducing
        )
        count = min(len(available_indices), needed)
        choices = torch.randperm(
            len(available_indices),
            generator=torch.Generator().manual_seed(self.seed),
        )[:count]
        indices = available_indices[choices]
        source = subset_loader(loader, indices)

    for batch in source:
        x, y = self._batch(batch)
        if self.likelihood == Likelihood.REWARD_MODELING:
            x = _individual_reward_inputs(x, self.dict_key_x)
            labels = torch.zeros(
                x[self.dict_key_x].shape[0]
                if isinstance(x, MutableMapping)
                else x.shape[0],
                device=self._device,
                dtype=torch.long,
            )
        elif self.likelihood == Likelihood.CLASSIFICATION:
            labels = classification_targets(y)
        else:
            labels = torch.zeros(y.shape[0], device=self._device, dtype=torch.long)
        candidates.append(x)
        classes.append(labels)

    if not candidates:
        raise ValueError("Cannot initialize inducing inputs from an empty loader.")

    inputs = concatenate_rows(candidates)
    labels = torch.cat(classes)
    available = labels.shape[0]
    if available < self.num_inducing:
        raise ValueError("num_inducing exceeds available training inputs.")
    if self._inducing_strategy == "kmeans":
        if isinstance(inputs, MutableMapping) or not inputs.is_floating_point():
            raise ValueError("kmeans requires floating-point tensor inputs.")
        flat = inputs.detach().cpu().reshape(available, -1).numpy()
        centers, _ = kmeans2(
            flat, self.num_inducing, minit="points", seed=self.seed
        )
        distances = ((flat[:, None, :] - centers[None, :, :]) ** 2).sum(-1)
        indices = torch.as_tensor(distances.argmin(0), device=self._device)
    else:
        indices = torch.randperm(
            available, generator=torch.Generator().manual_seed(self.seed)
        )[: self.num_inducing].to(self._device)

    self.inducing_locations = self._make_inducing(
        select_rows(inputs, indices, available)
    )
    self.inducing_classes = labels[indices].long()

_query_jacobians

_query_jacobians(
    x: Tensor | MutableMapping,
    *,
    differentiable: bool = False
) -> tuple[Tensor, Tensor]

Return model Jacobians and outputs for query inputs.

Parameters:

• x (Tensor or MutableMapping) –

  Batched inputs.

• differentiable (bool, default: False ) –

  Preserve gradients through Jacobian computation.

Returns:

• tuple[Tensor, Tensor] –

  Parameter Jacobians shaped [batch, outputs, parameters] and network outputs shaped [batch, outputs].

Source code in laplace/baselaplace.py

def _query_jacobians(
    self, x: torch.Tensor | MutableMapping, *, differentiable: bool = False
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return model Jacobians and outputs for query inputs.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched inputs.
    differentiable : bool, default=False
        Preserve gradients through Jacobian computation.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        Parameter Jacobians shaped ``[batch, outputs, parameters]`` and
        network outputs shaped ``[batch, outputs]``.
    """
    if self._uses_asdl_backend or (
        self.likelihood == Likelihood.REWARD_MODELING
        and not self._uses_backpack_backend
    ):
        return CurvatureInterface.jacobians(
            self.backend, x, enable_backprop=differentiable
        )
    output_size = self.model(x).shape[-1]
    previous = getattr(self.model, "output_size", None)
    setattr(self.model, "output_size", output_size)
    try:
        return self.backend.jacobians(x, enable_backprop=differentiable)
    finally:
        if previous is None:
            delattr(self.model, "output_size")
        else:
            setattr(self.model, "output_size", previous)

_inducing_jacobians

_inducing_jacobians() -> Tensor

Return selected-output Jacobians at all inducing locations.

Returns:

• Tensor –

  Parameter Jacobians shaped [num_inducing, parameters], with one output coordinate per location. Gradients through trainable inducing locations are retained.

Raises:

• RuntimeError –

  If inducing locations have not been initialized.

Source code in laplace/baselaplace.py

def _inducing_jacobians(self) -> torch.Tensor:
    """Return selected-output Jacobians at all inducing locations.

    Returns
    -------
    torch.Tensor
        Parameter Jacobians shaped ``[num_inducing, parameters]``, with
        one output coordinate per location. Gradients through trainable
        inducing locations are retained.

    Raises
    ------
    RuntimeError
        If inducing locations have not been initialized.
    """
    if self.inducing_locations is None:
        raise RuntimeError("Inducing locations have not been initialized.")
    selected, _ = self.backend.selected_output_jacobians(
        self.inducing_locations,
        self.inducing_classes,
        enable_backprop=isinstance(self.inducing_locations, torch.nn.Parameter),
    )
    return selected[:, 0, :]

_inducing_factor

_inducing_factor(
    jacobians: Tensor,
) -> tuple[Tensor, Tensor, Tensor]

Factor the inducing system without forming its Gram matrix.

Parameters:

• jacobians (Tensor) –

  Selected-output inducing Jacobians shaped [num_inducing, parameters].

Returns:

• features ( Tensor ) –

  Variationally transformed inducing Jacobians divided by the square root of prior precision, shaped [num_inducing, parameters].

• basis ( Tensor ) –

  Reduced QR basis of the vertically stacked feature transpose and identity, shaped [parameters + num_inducing, num_inducing].

• triangular ( Tensor ) –

  Upper-triangular QR factor shaped [num_inducing, num_inducing].

Notes

The factorization uses float64 arithmetic for numerical stability.

Source code in laplace/baselaplace.py

def _inducing_factor(
    self, jacobians: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Factor the inducing system without forming its Gram matrix.

    Parameters
    ----------
    jacobians : torch.Tensor
        Selected-output inducing Jacobians shaped
        ``[num_inducing, parameters]``.

    Returns
    -------
    features : torch.Tensor
        Variationally transformed inducing Jacobians divided by the
        square root of prior precision, shaped
        ``[num_inducing, parameters]``.
    basis : torch.Tensor
        Reduced QR basis of the vertically stacked feature transpose
        and identity, shaped ``[parameters + num_inducing, num_inducing]``.
    triangular : torch.Tensor
        Upper-triangular QR factor shaped ``[num_inducing, num_inducing]``.

    Notes
    -----
    The factorization uses float64 arithmetic for numerical stability.
    """
    work_dtype = torch.float64
    prior_scale = self.prior_precision.to(work_dtype).sqrt()
    inducing = jacobians.to(work_dtype)
    factor = self._factor_matrix().to(work_dtype)
    features = factor.T @ inducing / prior_scale
    identity = torch.eye(
        self.num_inducing, device=features.device, dtype=work_dtype
    )
    basis, triangular = torch.linalg.qr(
        torch.cat((features.T, identity), dim=0), mode="reduced"
    )
    return features, basis, triangular

_posterior_covariance

_posterior_covariance(
    jacobians: Tensor,
    inducing_jacobians: Tensor,
    *,
    joint: bool,
    basis: Tensor | None = None
) -> Tensor

Form latent covariance using the orthogonal inducing complement.

Parameters:

• jacobians (Tensor) –

  Query Jacobians shaped [batch, outputs, parameters].

• inducing_jacobians (Tensor) –

  Selected-output inducing Jacobians shaped [num_inducing, parameters].

• joint (bool) –

  Include covariance between query inputs when true.

• basis (Tensor or None, default: None ) –

  Precomputed orthonormal inducing basis shaped [parameters + num_inducing, num_inducing]. Computed from inducing_jacobians when omitted.

Returns:

• Tensor –

  Latent covariance shaped [batch, outputs, outputs], or [batch * outputs, batch * outputs] when joint=True, in the query Jacobian dtype. Observation noise is excluded.

Source code in laplace/baselaplace.py

def _posterior_covariance(
    self,
    jacobians: torch.Tensor,
    inducing_jacobians: torch.Tensor,
    *,
    joint: bool,
    basis: torch.Tensor | None = None,
) -> torch.Tensor:
    """Form latent covariance using the orthogonal inducing complement.

    Parameters
    ----------
    jacobians : torch.Tensor
        Query Jacobians shaped ``[batch, outputs, parameters]``.
    inducing_jacobians : torch.Tensor
        Selected-output inducing Jacobians shaped
        ``[num_inducing, parameters]``.
    joint : bool
        Include covariance between query inputs when true.
    basis : torch.Tensor or None, default=None
        Precomputed orthonormal inducing basis shaped
        ``[parameters + num_inducing, num_inducing]``. Computed from
        ``inducing_jacobians`` when omitted.

    Returns
    -------
    torch.Tensor
        Latent covariance shaped ``[batch, outputs, outputs]``, or
        ``[batch * outputs, batch * outputs]`` when ``joint=True``, in
        the query Jacobian dtype. Observation noise is excluded.
    """
    if basis is None:
        _, basis, _ = self._inducing_factor(inducing_jacobians)
    prior_scale = self.prior_precision.to(torch.float64).sqrt()
    query = jacobians.to(torch.float64).reshape(-1, self.n_params) / prior_scale
    padded = torch.nn.functional.pad(query, (0, self.num_inducing))
    roots = padded - (padded @ basis) @ basis.T
    if joint:
        covariance = roots @ roots.T
    else:
        batch, outputs = jacobians.shape[:2]
        roots = roots.reshape(batch, outputs, -1)
        covariance = torch.einsum("bcp,bdp->bcd", roots, roots)
    covariance = (covariance + covariance.transpose(-1, -2)) / 2
    return covariance.to(jacobians.dtype)

_latent_distribution

_latent_distribution(
    x: Tensor | MutableMapping, joint: bool = False
) -> tuple[Tensor, Tensor, Tensor, Tensor]

Return latent moments and factors for the variational objective.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• joint (bool, default: False ) –

  Retain covariance between input rows.

Returns:

• mean ( Tensor ) –

  Network outputs shaped [batch, outputs], including when joint=True.

• covariance ( Tensor ) –

  Latent covariance shaped [batch, outputs, outputs], or [batch * outputs, batch * outputs] when joint=True. Observation noise is excluded.

• features ( Tensor ) –

  Float64 scaled variational features shaped [num_inducing, parameters].

• triangular ( Tensor ) –

  Float64 upper-triangular inducing-system factor shaped [num_inducing, num_inducing].

Source code in laplace/baselaplace.py

def _latent_distribution(
    self, x: torch.Tensor | MutableMapping, joint: bool = False
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Return latent moments and factors for the variational objective.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    joint : bool, default=False
        Retain covariance between input rows.

    Returns
    -------
    mean : torch.Tensor
        Network outputs shaped ``[batch, outputs]``, including when
        ``joint=True``.
    covariance : torch.Tensor
        Latent covariance shaped ``[batch, outputs, outputs]``, or
        ``[batch * outputs, batch * outputs]`` when ``joint=True``.
        Observation noise is excluded.
    features : torch.Tensor
        Float64 scaled variational features shaped
        ``[num_inducing, parameters]``.
    triangular : torch.Tensor
        Float64 upper-triangular inducing-system factor shaped
        ``[num_inducing, num_inducing]``.
    """
    x = to_device(x, self._device, self._dtype)
    jacobians, mean = self._query_jacobians(x, differentiable=self.enable_backprop)
    inducing_jacobians = self._inducing_jacobians()
    features, basis, triangular = self._inducing_factor(inducing_jacobians)
    covariance = self._posterior_covariance(
        jacobians, inducing_jacobians, joint=joint, basis=basis
    )
    return mean, covariance, features, triangular

_glm_predictive_distribution

_glm_predictive_distribution(
    x: Tensor | MutableMapping, joint: bool = False
) -> tuple[Tensor, Tensor]

Compute model outputs and VaLLA latent covariance for inputs.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• joint (bool, default: False ) –

  Retain covariance between input rows.

Returns:

• mean ( Tensor ) –

  Network outputs shaped [batch, outputs], flattened to [batch * outputs] when joint=True.

• covariance ( Tensor ) –

  Latent covariance shaped [batch, outputs, outputs], or [batch * outputs, batch * outputs] when joint=True. Observation noise is excluded.

Source code in laplace/baselaplace.py

@preserve_model_gradients
@torch.enable_grad()
def _glm_predictive_distribution(
    self, x: torch.Tensor | MutableMapping, joint: bool = False
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute model outputs and VaLLA latent covariance for inputs.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    joint : bool, default=False
        Retain covariance between input rows.

    Returns
    -------
    mean : torch.Tensor
        Network outputs shaped ``[batch, outputs]``, flattened to
        ``[batch * outputs]`` when ``joint=True``.
    covariance : torch.Tensor
        Latent covariance shaped ``[batch, outputs, outputs]``, or
        ``[batch * outputs, batch * outputs]`` when ``joint=True``.
        Observation noise is excluded.
    """
    mean, covariance, _, _ = self._latent_distribution(x, joint)
    if joint:
        mean = mean.flatten()
    if not self.enable_backprop:
        return mean.detach(), covariance.detach()
    return mean, covariance

functional_variance

functional_variance(jacobians: Tensor) -> Tensor

Return per-input latent covariance from parameter Jacobians.

Parameters:

• jacobians (Tensor) –

  Parameter Jacobians shaped [batch, outputs, parameters].

Returns:

• Tensor –

  Per-input covariance shaped [batch, outputs, outputs].

Source code in laplace/baselaplace.py

@preserve_model_gradients
def functional_variance(self, jacobians: torch.Tensor) -> torch.Tensor:
    """Return per-input latent covariance from parameter Jacobians.

    Parameters
    ----------
    jacobians : torch.Tensor
        Parameter Jacobians shaped ``[batch, outputs, parameters]``.

    Returns
    -------
    torch.Tensor
        Per-input covariance shaped ``[batch, outputs, outputs]``.
    """
    self._check_fitted()
    self._check_jacobians(jacobians)
    inducing = self._inducing_jacobians()
    return self._posterior_covariance(jacobians, inducing, joint=False)

functional_covariance

functional_covariance(jacobians: Tensor) -> Tensor

Return joint latent covariance from parameter Jacobians.

Parameters:

• jacobians (Tensor) –

  Parameter Jacobians shaped [batch, outputs, parameters].

Returns:

• Tensor –

  Joint covariance shaped [batch * outputs, batch * outputs].

Source code in laplace/baselaplace.py

@preserve_model_gradients
def functional_covariance(self, jacobians: torch.Tensor) -> torch.Tensor:
    """Return joint latent covariance from parameter Jacobians.

    Parameters
    ----------
    jacobians : torch.Tensor
        Parameter Jacobians shaped ``[batch, outputs, parameters]``.

    Returns
    -------
    torch.Tensor
        Joint covariance shaped ``[batch * outputs, batch * outputs]``.
    """
    self._check_fitted()
    self._check_jacobians(jacobians)
    inducing = self._inducing_jacobians()
    return self._posterior_covariance(jacobians, inducing, joint=True)

_objective

_objective(x: Any, y: Tensor) -> Tensor

Return the minibatch-scaled alpha objective plus inducing KL.

Parameters:

• x (Any) –

  Batched model inputs.

• y (Tensor) –

  Regression targets or class labels.

Returns:

• Tensor –

  Scalar loss minimized during fitting.

Source code in laplace/baselaplace.py

def _objective(self, x: Any, y: torch.Tensor) -> torch.Tensor:
    """Return the minibatch-scaled alpha objective plus inducing KL.

    Parameters
    ----------
    x : Any
        Batched model inputs.
    y : torch.Tensor
        Regression targets or class labels.

    Returns
    -------
    torch.Tensor
        Scalar loss minimized during fitting.
    """
    mean, covariance, features, triangular = self._latent_distribution(x)
    if self.likelihood == Likelihood.REGRESSION:
        y = regression_targets(y, mean)
        noise_variance = self.sigma_noise.square()
        latent_variance = covariance.diagonal(dim1=-2, dim2=-1)
        squared_error = (y - mean).square()
        if self.alpha == 0:
            energy = (squared_error + latent_variance) / noise_variance
        else:
            energy = (
                squared_error / (noise_variance + self.alpha * latent_variance)
                + torch.log1p(self.alpha * latent_variance / noise_variance)
                / self.alpha
            )
        log_term = -0.5 * (energy + torch.log(2 * torch.pi * noise_variance)).sum()
    else:
        y = classification_targets(y)
        if self.mc_softmax_samples:
            draws = draw_gaussian(
                mean, covariance, self.mc_softmax_samples, self.generator
            )
            log_probabilities = draws.log_softmax(dim=-1)
            selected = log_probabilities.gather(
                -1,
                y.long().view(1, -1, 1).expand(self.mc_softmax_samples, -1, 1),
            ).squeeze(-1)
            if self.alpha == 0:
                log_term = selected.mean(dim=0).sum()
            else:
                log_term = (
                    torch.logsumexp(self.alpha * selected, dim=0)
                    - torch.log(
                        torch.as_tensor(
                            self.mc_softmax_samples,
                            device=self._device,
                            dtype=self._dtype,
                        )
                    )
                ).sum() / self.alpha
        else:
            scaled = mean / torch.sqrt(
                1 + torch.pi / 8 * covariance.diagonal(dim1=-2, dim2=-1)
            )
            log_term = -torch.nn.functional.cross_entropy(
                scaled, y.long(), reduction="sum"
            )

    solved = torch.linalg.solve_triangular(triangular.T, features, upper=False)
    kl = triangular.diagonal().abs().log().sum() - 0.5 * solved.square().sum()
    return -(self.n_data / y.shape[0]) * log_term / self.temperature + kl

_validation_nll

_validation_nll(loader: DataLoader) -> Tensor

Return mean predictive negative log likelihood on a loader.

Parameters:

• loader (DataLoader) –

  Nonempty validation loader with regression targets or class labels.

Returns:

• Tensor –

  Scalar NLL averaged over examples, summing over regression output coordinates. Regression includes observation noise; classification uses the probit approximation.

Source code in laplace/baselaplace.py

def _validation_nll(self, loader: DataLoader) -> torch.Tensor:
    """Return mean predictive negative log likelihood on a loader.

    Parameters
    ----------
    loader : DataLoader
        Nonempty validation loader with regression targets or class labels.

    Returns
    -------
    torch.Tensor
        Scalar NLL averaged over examples, summing over regression output
        coordinates. Regression includes observation noise; classification
        uses the probit approximation.
    """
    total = torch.zeros((), device=self._device, dtype=self._dtype)
    count = 0
    for batch in loader:
        x, y = self._batch(batch)
        if self.likelihood == Likelihood.REGRESSION:
            mean, covariance = self.predictive_moments(x)
            y = regression_targets(y, mean)
            variance = (
                covariance.diagonal(dim1=-2, dim2=-1) + self.sigma_noise.square()
            )
            total += (
                0.5
                * (
                    (y - mean).square() / variance
                    + variance.log()
                    + torch.log(torch.as_tensor(2 * torch.pi, device=self._device))
                )
            ).sum()
        else:
            total += self._classification_validation_nll(x, y)
        count += y.shape[0]
    return total / count

fit

fit(
    train_loader: DataLoader,
    *,
    iterations: int = 100,
    lr: float = 0.001,
    val_loader: DataLoader | None = None,
    val_steps: int | None = None,
    progress_bar: bool = False,
    override: bool = True
) -> None

Optimize the VaLLA variational posterior and hyperparameters.

Parameters:

• train_loader (DataLoader) –

  Training batches used for inducing selection and optimization.

• iterations (int, default: 100 ) –

  Number of minibatch optimizer steps.

• lr (float, default: 1e-3 ) –

  Adam learning rate.

• val_loader (DataLoader or None, default: None ) –

  Optional data for predictive negative log likelihood.

• val_steps (int or None, default: None ) –

  Evaluate validation NLL every val_steps iterations when given.

• progress_bar (bool, default: False ) –

  Show an optimization progress bar.

• override (bool, default: True ) –

  Reset inducing and variational state; false continues a prior fit.

Raises:

• ValueError –

  If optimization settings or inducing classes are invalid.

Source code in laplace/baselaplace.py

@preserve_model_gradients
def fit(
    self,
    train_loader: DataLoader,
    *,
    iterations: int = 100,
    lr: float = 1e-3,
    val_loader: DataLoader | None = None,
    val_steps: int | None = None,
    progress_bar: bool = False,
    override: bool = True,
) -> None:
    """Optimize the VaLLA variational posterior and hyperparameters.

    Parameters
    ----------
    train_loader : DataLoader
        Training batches used for inducing selection and optimization.
    iterations : int, default=100
        Number of minibatch optimizer steps.
    lr : float, default=1e-3
        Adam learning rate.
    val_loader : DataLoader or None, default=None
        Optional data for predictive negative log likelihood.
    val_steps : int or None, default=None
        Evaluate validation NLL every ``val_steps`` iterations when given.
    progress_bar : bool, default=False
        Show an optimization progress bar.
    override : bool, default=True
        Reset inducing and variational state; false continues a prior fit.

    Raises
    ------
    ValueError
        If optimization settings or inducing classes are invalid.
    """
    if (
        not isinstance(iterations, Integral)
        or isinstance(iterations, bool)
        or iterations <= 0
        or not isfinite(lr)
        or lr <= 0
    ):
        raise ValueError(
            "iterations must be a positive integer and lr positive and finite."
        )
    if val_steps is not None and (
        not isinstance(val_steps, Integral)
        or isinstance(val_steps, bool)
        or val_steps <= 0
    ):
        raise ValueError("val_steps must be positive integer when provided.")
    if not isfinite(self.temperature) or self.temperature <= 0:
        raise ValueError("temperature must be positive and finite.")
    if not 0 <= self.alpha <= 1:
        raise ValueError("alpha must be in [0, 1].")
    if not isinstance(self.mc_softmax_samples, int) or self.mc_softmax_samples < 0:
        raise ValueError("mc_softmax_samples must be a nonnegative integer.")
    if (
        self.likelihood != Likelihood.REGRESSION
        and self.alpha != 1
        and self.mc_softmax_samples == 0
    ):
        raise ValueError(
            "Classification alpha != 1 requires mc_softmax_samples > 0."
        )

    self.model.eval()
    self.n_data = len(training_indices(train_loader))
    x_first, _ = self._batch(next(iter(train_loader)))
    if isinstance(x_first, MutableMapping) and self._uses_backpack_backend:
        raise ValueError(
            "BackPACK does not support mapping-style inputs; use AsdlGGN or CurvlinopsGGN."
        )

    self.n_outputs = self.model(x_first).shape[-1]
    setattr(self.model, "output_size", self.n_outputs)
    if override or not self._fitted:
        if self._inducing_strategy is not None:
            self._gather_inducing(train_loader)
        else:
            self.inducing_locations = self._make_inducing(
                self._initial_inducing_locations
            )
            self.inducing_classes = self._initial_inducing_classes.detach().clone()
        self._init_variational_factor()
        self.fit_history_ = {"objective": [], "val_nll": []}

    if (
        self.inducing_classes.shape != (self.num_inducing,)
        or torch.any(self.inducing_classes < 0)
        or torch.any(self.inducing_classes >= self.n_outputs)
    ):
        raise ValueError(
            "inducing_classes must contain one valid output per inducing input."
        )

    self._fitted = False
    parameters = [self.L, self.log_prior_precision]
    if self.log_noise_variance is not None:
        parameters.append(self.log_noise_variance)
    if isinstance(self.inducing_locations, torch.nn.Parameter):
        parameters.append(self.inducing_locations)

    optimizer = torch.optim.Adam(parameters, lr=lr)
    iterator = range(iterations)
    if progress_bar:
        from tqdm import trange

        iterator = trange(iterations, desc="Fitting VaLLA")

    batches = iter(train_loader)
    for step in iterator:
        try:
            batch = next(batches)
        except StopIteration:
            batches = iter(train_loader)
            batch = next(batches)
        x, y = self._batch(batch)
        optimizer.zero_grad()
        objective = self._objective(x, y)
        objective.backward()
        optimizer.step()
        self.fit_history_["objective"].append(float(objective.detach()))
        if val_loader is not None and val_steps and (step + 1) % val_steps == 0:
            self._fitted = True
            with torch.no_grad():
                self.fit_history_["val_nll"].append(
                    float(self._validation_nll(val_loader))
                )

    self._fitted = True
    if val_loader is not None and (not val_steps or iterations % val_steps):
        with torch.no_grad():
            self.fit_history_["val_nll"].append(
                float(self._validation_nll(val_loader))
            )

optimize_prior_precision

optimize_prior_precision(*args, **kwargs) -> None

Reject separate prior tuning because :meth:fit learns the precision.

Parameters:

• *args (Any, default: () ) –

  Positional arguments accepted for compatibility with BaseLaplace.

• **kwargs (Any, default: {} ) –

  Keyword arguments accepted for compatibility with BaseLaplace.

Raises:

• NotImplementedError –

  Always raised. VaLLA optimizes prior precision jointly with the variational parameters during fitting.

Source code in laplace/baselaplace.py

def optimize_prior_precision(self, *args, **kwargs) -> None:
    """Reject separate prior tuning because :meth:`fit` learns the precision.

    Parameters
    ----------
    *args : Any
        Positional arguments accepted for compatibility with BaseLaplace.
    **kwargs : Any
        Keyword arguments accepted for compatibility with BaseLaplace.

    Raises
    ------
    NotImplementedError
        Always raised. VaLLA optimizes prior precision jointly with the
        variational parameters during fitting.
    """
    raise NotImplementedError("VaLLA optimizes prior precision during fit.")

log_marginal_likelihood

log_marginal_likelihood(*args, **kwargs)

Reject marginal-likelihood evaluation for the VaLLA approximation.

Parameters:

• *args (Any, default: () ) –

  Positional arguments accepted for compatibility with BaseLaplace.

• **kwargs (Any, default: {} ) –

  Keyword arguments accepted for compatibility with BaseLaplace.

Raises:

• NotImplementedError –

  Always raised. The optimized alpha-divergence objective is not a Laplace marginal likelihood.

Source code in laplace/baselaplace.py

def log_marginal_likelihood(self, *args, **kwargs):
    """Reject marginal-likelihood evaluation for the VaLLA approximation.

    Parameters
    ----------
    *args : Any
        Positional arguments accepted for compatibility with BaseLaplace.
    **kwargs : Any
        Keyword arguments accepted for compatibility with BaseLaplace.

    Raises
    ------
    NotImplementedError
        Always raised. The optimized alpha-divergence objective is not a
        Laplace marginal likelihood.
    """
    raise NotImplementedError(
        "VaLLA's variational objective is not a Laplace marginal likelihood."
    )

state_dict

state_dict() -> dict[str, Any]

Return a checkpoint containing fitted variational and inducing state.

Returns:

• dict[str, Any] –

  Variational factor, inducing locations and output coordinates, hyperparameters, fitting history, initialization settings, random generator state, model fingerprint, and backend/input settings.

Raises:

• RuntimeError –

  If the posterior has not been fitted.

Notes

Loading validates the pretrained model against the saved fingerprint; it does not restore model parameters or buffers. Construct the receiving estimator with the same pretrained model and configuration.

Source code in laplace/baselaplace.py

def state_dict(self) -> dict[str, Any]:
    """Return a checkpoint containing fitted variational and inducing state.

    Returns
    -------
    dict[str, Any]
        Variational factor, inducing locations and output coordinates,
        hyperparameters, fitting history, initialization settings, random
        generator state, model fingerprint, and backend/input settings.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.

    Notes
    -----
    Loading validates the pretrained model against the saved fingerprint;
    it does not restore model parameters or buffers. Construct the
    receiving estimator with the same pretrained model and configuration.
    """
    self._check_fitted()
    return {
        "cls_name": type(self).__name__,
        "n_params": self.n_params,
        "model_fingerprint": _model_fingerprint(self.model),
        "likelihood": self.likelihood,
        "inducing_locations": self.inducing_locations.detach().clone()
        if isinstance(self.inducing_locations, torch.Tensor)
        else deepcopy(self.inducing_locations),
        "inducing_classes": self.inducing_classes.detach().clone(),
        "inducing_strategy": self._inducing_strategy,
        "initial_inducing_locations": self._clone_fixed_inputs(
            self._initial_inducing_locations
        ),
        "initial_inducing_classes": self._initial_inducing_classes.detach().clone(),
        "L": self.L.detach().clone(),
        "log_prior_precision": self.log_prior_precision.detach().clone(),
        "log_noise_variance": None
        if self.log_noise_variance is None
        else self.log_noise_variance.detach().clone(),
        "temperature": self.temperature,
        "alpha": self.alpha,
        "mc_softmax_samples": self.mc_softmax_samples,
        "n_data": self.n_data,
        "n_outputs": self.n_outputs,
        "fitted": self._fitted,
        "fit_history": deepcopy(self.fit_history_),
        "seed": self.seed,
        "generator_state": self.generator.get_state().clone(),
        **self._checkpoint_configuration(),
    }

load_state_dict

load_state_dict(state_dict: dict[str, Any]) -> None

Restore a fitted VaLLA checkpoint for the same pretrained model.

Parameters:

• state_dict (dict[str, Any]) –

  State previously returned by :meth:state_dict.

Raises:

• ValueError –

  If the checkpoint is unfitted or incompatible with this model.

Source code in laplace/baselaplace.py

def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Restore a fitted VaLLA checkpoint for the same pretrained model.

    Parameters
    ----------
    state_dict : dict[str, Any]
        State previously returned by :meth:`state_dict`.

    Raises
    ------
    ValueError
        If the checkpoint is unfitted or incompatible with this model.
    """
    if not state_dict.get("fitted", False):
        raise ValueError("VaLLA checkpoint must contain a fitted posterior.")
    if (
        state_dict["cls_name"] != type(self).__name__
        or state_dict["n_params"] != self.n_params
    ):
        raise ValueError(
            "Checkpoint requires the same VaLLA type and pretrained model."
        )
    if state_dict["likelihood"] != self.likelihood:
        raise ValueError("Checkpoint likelihood does not match.")
    self._validate_checkpoint_configuration(state_dict)
    _check_model_fingerprint(self.model, state_dict["model_fingerprint"])
    self.seed = state_dict["seed"]
    self.generator.set_state(state_dict["generator_state"])
    self.inducing_locations = self._make_inducing(state_dict["inducing_locations"])
    self.inducing_classes = state_dict["inducing_classes"].to(self._device)
    self.num_inducing = len(self.inducing_classes)
    self._inducing_strategy = state_dict["inducing_strategy"]
    self._initial_inducing_locations = self._clone_fixed_inputs(
        to_device(
            state_dict["initial_inducing_locations"], self._device, self._dtype
        )
    )
    self._initial_inducing_classes = state_dict["initial_inducing_classes"].to(
        self._device
    )
    self._init_variational_factor()
    with torch.no_grad():
        self.L.copy_(state_dict["L"].to(self._device))
        self.log_prior_precision.copy_(
            state_dict["log_prior_precision"].to(self._device)
        )
        if self.log_noise_variance is not None:
            self.log_noise_variance.copy_(
                state_dict["log_noise_variance"].to(self._device)
            )
    self.n_data = state_dict["n_data"]
    self.n_outputs = state_dict["n_outputs"]
    self.temperature = state_dict["temperature"]
    self.alpha = state_dict["alpha"]
    self.mc_softmax_samples = state_dict["mc_softmax_samples"]
    setattr(self.model, "output_size", self.n_outputs)
    self._fitted = state_dict["fitted"]
    self.fit_history_ = deepcopy(state_dict["fit_history"])
    self.model.eval()