ELLA

laplace.baselaplace.ELLA

ELLA(
    model: Module,
    likelihood: Likelihood | str,
    subsample_size: int,
    n_eigenvalues: int,
    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 = 1234,
)

Bases: BaseFunctionalLaplace

Accelerated linearized Laplace using a rank-limited Nyström feature map.

The pretrained model stays fixed. Fitting selects a subset of training inputs, builds parameter-Jacobian features, and accumulates projected GGN curvature. A numerically rank-deficient subset cannot supply the requested number of features.

Parameters:

• model (Module) –

  Pretrained network to linearize.

• likelihood (Likelihood or str) –

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

• subsample_size (int) –

  Number of training inputs used to construct the Nyström basis.

• n_eigenvalues (int) –

  Number of positive kernel directions retained, at most subsample_size.

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

  Observation standard deviation for regression.

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

  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 curvature.

• 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 to compute parameter Jacobians.

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

  Additional curvature backend options.

• seed (int, default: 1234 ) –

  Seed for selecting basis inputs and output coordinates.

Attributes:

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

  processed_examples records cumulative curvature-pass counts, val_nll records validation scores, and tuning records candidate hyperparameters and their validation scores.

• dual_directions (Tensor or None) –

  Fitted parameter-space basis shaped [parameters, n_eigenvalues].

• feature_gram (Tensor or None) –

  Accumulated projected curvature shaped [n_eigenvalues, n_eigenvalues] before noise and temperature scaling.

• chol_precision (Tensor or None) –

  Lower Cholesky factor of the feature-space posterior precision.

Notes

fit returns None. fit_history_ stores processed example counts, validation NLL, and prior-grid scores when those options are used.

Source code in laplace/baselaplace.py

def __init__(
    self,
    model: torch.nn.Module,
    likelihood: Likelihood | str,
    subsample_size: int,
    n_eigenvalues: int,
    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 = 1234,
) -> None:
    """Initialize ELLA from the pretrained model and basis configuration."""
    if (
        not isinstance(subsample_size, Integral)
        or isinstance(subsample_size, bool)
        or not isinstance(n_eigenvalues, Integral)
        or isinstance(n_eigenvalues, bool)
        or subsample_size <= 0
        or n_eigenvalues <= 0
        or n_eigenvalues > subsample_size
    ):
        raise ValueError("Require 0 < n_eigenvalues <= subsample_size.")
    if torch.as_tensor(prior_precision).numel() != 1:
        raise ValueError("ELLA requires scalar prior_precision.")
    if torch.any(torch.as_tensor(prior_mean) != 0):
        raise ValueError("ELLA currently requires prior_mean=0.")
    if not isfinite(temperature) or temperature <= 0:
        raise ValueError("temperature must be positive and finite.")
    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,
    )
    self.subsample_size = subsample_size
    self.n_eigenvalues = n_eigenvalues
    self.seed = seed
    self.dual_directions: torch.Tensor | None = None
    self.feature_gram: torch.Tensor | None = None
    self.chol_precision: torch.Tensor | None = None
    self._posterior_dirty = False
    self._fitted = False
    self.fit_history_: dict[str, list[Any]] = {
        "processed_examples": [],
        "val_nll": [],
        "tuning": [],
    }

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

n_data

n_data: 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.

subsample_size

subsample_size = subsample_size

n_eigenvalues

n_eigenvalues = n_eigenvalues

seed

seed = seed

dual_directions

dual_directions: Tensor | None = None

feature_gram

feature_gram: Tensor | None = None

chol_precision

chol_precision: Tensor | None = None

_posterior_dirty

_posterior_dirty = False

_fitted

_fitted = False

fit_history_

fit_history_: dict[str, list[Any]] = {
    "processed_examples": [],
    "val_nll": [],
    "tuning": [],
}

prior_precision

prior_precision: Tensor

Return the scalar parameter-prior precision.

prior_mean

prior_mean: Tensor

Return the zero parameter-prior mean.

sigma_noise

sigma_noise: Tensor

Return the regression observation standard deviation.

temperature

temperature: float

Return the positive likelihood temperature.

log_likelihood

log_likelihood: Tensor

Reject access to an untracked training log likelihood.

Raises:

• NotImplementedError –

  Always raised because ELLA accumulates curvature without tracking the 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}.")

_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

_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)

_indices

_indices(loader: DataLoader, balanced: bool) -> Tensor

Select reproducible basis rows, optionally balanced across classes.

Parameters:

• loader (DataLoader) –

  Training data from which basis rows are selected.

• balanced (bool) –

  Round-robin sample classes when true.

Returns:

• Tensor –

  Selected dataset indices of length subsample_size.

Source code in laplace/baselaplace.py

def _indices(self, loader: DataLoader, balanced: bool) -> torch.Tensor:
    """Select reproducible basis rows, optionally balanced across classes.

    Parameters
    ----------
    loader : DataLoader
        Training data from which basis rows are selected.
    balanced : bool
        Round-robin sample classes when true.

    Returns
    -------
    torch.Tensor
        Selected dataset indices of length ``subsample_size``.
    """
    available = training_indices(loader)
    n = len(available)
    if self.subsample_size > n:
        raise ValueError("subsample_size cannot exceed the training set size.")

    generator = torch.Generator().manual_seed(self.seed)
    if not balanced:
        chosen = torch.randperm(n, generator=generator)[: self.subsample_size]
        return available[chosen]

    if self.likelihood == Likelihood.REGRESSION:
        raise ValueError("balanced sampling is only defined for classification.")

    labels = []
    for index in available:
        _, target = split_batch(loader.dataset[int(index)], self.dict_key_y)
        labels.append(
            int(
                torch.as_tensor(target).argmax()
                if torch.as_tensor(target).numel() > 1
                else target
            )
        )
    labels = torch.tensor(labels)
    classes = labels.unique(sorted=True)
    queues = []
    for cls in classes:
        candidates = torch.where(labels == cls)[0]
        order = torch.randperm(len(candidates), generator=generator)
        queues.append(deque(candidates[order].tolist()))

    selected: list[int] = []
    while len(selected) < self.subsample_size and any(queues):
        for queue in queues:
            if queue and len(selected) < self.subsample_size:
                selected.append(queue.popleft())
    if len(selected) != self.subsample_size:
        raise ValueError("Not enough examples for balanced subsampling.")
    return available[torch.tensor(selected[: self.subsample_size])]

_build_basis

_build_basis(loader: DataLoader, indices: Tensor) -> None

Build the Nyström directions from selected parameter Jacobians.

Parameters:

• loader (DataLoader) –

  Training loader used to retrieve selected inputs.

• indices (Tensor) –

  Dataset indices selected for the subset kernel.

Raises:

• ValueError –

  If the subset kernel has fewer usable directions than requested.

Source code in laplace/baselaplace.py

def _build_basis(self, loader: DataLoader, indices: torch.Tensor) -> None:
    """Build the Nyström directions from selected parameter Jacobians.

    Parameters
    ----------
    loader : DataLoader
        Training loader used to retrieve selected inputs.
    indices : torch.Tensor
        Dataset indices selected for the subset kernel.

    Raises
    ------
    ValueError
        If the subset kernel has fewer usable directions than requested.
    """
    selected_rows = []
    generator = torch.Generator().manual_seed(self.seed + 1)
    for batch in subset_loader(loader, indices):
        x, _ = self._batch(batch)
        if (
            self.likelihood == Likelihood.REWARD_MODELING
            and not self._uses_backpack_backend
        ):
            jacobians, outputs = CurvatureInterface.jacobians(self.backend, x)
        else:
            jacobians, outputs = self.backend.jacobians(x)
        choices = torch.randint(
            outputs.shape[-1], (outputs.shape[0],), generator=generator
        ).to(jacobians.device)
        selected_rows.append(
            jacobians[torch.arange(len(choices), device=jacobians.device), choices]
        )

    rows = torch.cat(selected_rows).to(torch.float64)
    kernel = rows @ rows.T
    eigenvalues, eigenvectors = torch.linalg.eigh(kernel)
    order = torch.argsort(eigenvalues, descending=True)[: self.n_eigenvalues]
    values = eigenvalues[order]
    tolerance = (
        torch.finfo(values.dtype).eps * max(rows.shape) * eigenvalues.abs().max()
    )

    if torch.any(values <= tolerance):
        usable_rank = int((eigenvalues > tolerance).sum().item())
        raise ValueError(
            "Nyström subset kernel has insufficient positive rank: "
            f"requested n_eigenvalues={self.n_eigenvalues}, but only "
            f"{usable_rank} of {len(eigenvalues)} sampled directions "
            "exceed the numerical tolerance. Reduce n_eigenvalues, "
            "increase subsample_size, or try a different seed."
        )

    self.dual_directions = (
        rows.T @ (eigenvectors[:, order] / values.sqrt().unsqueeze(0))
    ).to(self._dtype)

_features

_features(
    x: Tensor | MutableMapping,
) -> tuple[Tensor, Tensor]

Project model Jacobians onto the Nyström basis and return outputs.

Parameters:

• x (Tensor or MutableMapping) –

  Batched inputs on the model device.

Returns:

• features ( Tensor ) –

  Projected Jacobians shaped [batch, outputs, n_eigenvalues].

• outputs ( Tensor ) –

  Network outputs shaped [batch, outputs].

Raises:

• RuntimeError –

  If the basis is uninitialized, or differentiable Jacobians are requested but unavailable for the model and backend.

Notes

Uses Jacobian-vector products when supported, otherwise projects parameter Jacobians obtained from the curvature backend.

Source code in laplace/baselaplace.py

def _features(
    self, x: torch.Tensor | MutableMapping
) -> tuple[torch.Tensor, torch.Tensor]:
    """Project model Jacobians onto the Nyström basis and return outputs.

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

    Returns
    -------
    features : torch.Tensor
        Projected Jacobians shaped ``[batch, outputs, n_eigenvalues]``.
    outputs : torch.Tensor
        Network outputs shaped ``[batch, outputs]``.

    Raises
    ------
    RuntimeError
        If the basis is uninitialized, or differentiable Jacobians are
        requested but unavailable for the model and backend.

    Notes
    -----
    Uses Jacobian-vector products when supported, otherwise projects
    parameter Jacobians obtained from the curvature backend.
    """
    if self.dual_directions is None:
        raise RuntimeError("Call fit before requesting ELLA features.")
    # BackPACK extends modules in-place and is incompatible with torch.func.jvp.
    if not self._uses_backpack_backend:
        parameters = dict(self.model.named_parameters())
        buffers = dict(self.model.named_buffers())
        feature_columns = []
        for direction in self.dual_directions.T:
            offset = 0
            tangents = {}
            for name, param in parameters.items():
                if param.requires_grad:
                    tangent = direction[offset : offset + param.numel()].reshape_as(
                        param
                    )
                    offset += param.numel()
                else:
                    tangent = torch.zeros_like(param)
                tangents[name] = tangent

            def model_fn(params):
                return torch.func.functional_call(
                    self.model, (params, buffers), (x,)
                )

            try:
                output, product = torch.func.jvp(
                    model_fn, (parameters,), (tangents,)
                )
            except (NotImplementedError, RuntimeError) as exc:
                if "forward AD" not in str(exc) and "jvp" not in str(exc):
                    raise
                break
            feature_columns.append(product)
        if len(feature_columns) == self.n_eigenvalues:
            return torch.stack(feature_columns, dim=-1), output
    if (
        self.likelihood == Likelihood.REWARD_MODELING
        and self._uses_backpack_backend
        and (query_shape := self.model(x).shape)[-1] == 1
    ):
        selectors = torch.zeros(
            query_shape[0], device=self._device, dtype=torch.long
        )
        jacobians, output = self.backend.selected_output_jacobians(
            x, selectors, enable_backprop=self.enable_backprop
        )
    elif (
        self.likelihood == Likelihood.REWARD_MODELING
        and not self._uses_backpack_backend
    ):
        jacobians, output = CurvatureInterface.jacobians(
            self.backend, x, enable_backprop=self.enable_backprop
        )
    elif self._uses_asdl_backend:
        # Reverse-mode torch.func can retain input derivatives when the
        # model does not support JVP; ASDL's Jacobians can detach them.
        try:
            jacobians, output = CurvatureInterface.jacobians(
                self.backend, x, enable_backprop=self.enable_backprop
            )
        except (NotImplementedError, RuntimeError) as exc:
            if self.enable_backprop:
                raise RuntimeError(
                    "This model cannot provide differentiable ELLA Jacobians."
                ) from exc
            jacobians, output = self.backend.jacobians(x)
    else:
        jacobians, output = self.backend.jacobians(
            x, enable_backprop=self.enable_backprop
        )
    return jacobians @ self.dual_directions, output

_refresh_posterior

_refresh_posterior() -> None

Rebuild the posterior factor after a hyperparameter change.

Source code in laplace/baselaplace.py

def _refresh_posterior(self) -> None:
    """Rebuild the posterior factor after a hyperparameter change."""
    if self._posterior_dirty:
        self._build_precision()

_build_precision

_build_precision() -> None

Factor the feature-space posterior precision.

Source code in laplace/baselaplace.py

def _build_precision(self) -> None:
    """Factor the feature-space posterior precision."""
    if self.feature_gram is None:
        return
    precision = self.feature_gram / self.temperature
    if self.likelihood == Likelihood.REGRESSION:
        precision = precision / self.sigma_noise.square()
    precision = precision + self.prior_precision * torch.eye(
        self.n_eigenvalues, device=self._device, dtype=self._dtype
    )
    self.chol_precision = torch.linalg.cholesky(precision)
    self._posterior_dirty = False

fit

fit(
    train_loader: DataLoader,
    *,
    val_loader: DataLoader | None = None,
    val_steps: int | None = None,
    balanced: bool = False,
    progress_bar: bool = False
) -> None

Fit the ELLA posterior with projected training curvature.

Parameters:

• train_loader (DataLoader) –

  Training batches for basis selection and curvature accumulation.

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

  Optional validation data for predictive negative log likelihood.

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

  Evaluate validation NLL every val_steps batches when given.

• balanced (bool, default: False ) –

  Balance class labels in the Nyström subset.

• progress_bar (bool, default: False ) –

  Show a fitting progress bar.

Raises:

• ValueError –

  If sampling or validation-step settings are invalid, or the curvature pass processes no examples.

Notes

Basis selection and curvature accumulation sample the loader independently. Stochastic samplers may select different rows for each pass. n_data counts examples processed during curvature accumulation.

Source code in laplace/baselaplace.py

@preserve_model_gradients
def fit(
    self,
    train_loader: DataLoader,
    *,
    val_loader: DataLoader | None = None,
    val_steps: int | None = None,
    balanced: bool = False,
    progress_bar: bool = False,
) -> None:
    """Fit the ELLA posterior with projected training curvature.

    Parameters
    ----------
    train_loader : DataLoader
        Training batches for basis selection and curvature accumulation.
    val_loader : DataLoader or None, default=None
        Optional validation data for predictive negative log likelihood.
    val_steps : int or None, default=None
        Evaluate validation NLL every ``val_steps`` batches when given.
    balanced : bool, default=False
        Balance class labels in the Nyström subset.
    progress_bar : bool, default=False
        Show a fitting progress bar.

    Raises
    ------
    ValueError
        If sampling or validation-step settings are invalid, or the
        curvature pass processes no examples.

    Notes
    -----
    Basis selection and curvature accumulation sample the loader
    independently. Stochastic samplers may select different rows for each
    pass. ``n_data`` counts examples processed during curvature accumulation.
    """
    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.")

    self.model.eval()
    self.fit_history_ = {"processed_examples": [], "val_nll": [], "tuning": []}
    self._fitted = False
    self.n_data = 0
    indices = self._indices(train_loader, balanced)
    first_x, _ = self._batch(next(iter(subset_loader(train_loader, indices))))
    if isinstance(first_x, MutableMapping) and self._uses_backpack_backend:
        raise ValueError(
            "BackPACK does not support mapping-style inputs; use AsdlGGN or CurvlinopsGGN."
        )

    self.n_outputs = self.model(first_x).shape[-1]
    setattr(self.model, "output_size", self.n_outputs)
    self._build_basis(train_loader, indices)
    self.feature_gram = torch.zeros(
        self.n_eigenvalues,
        self.n_eigenvalues,
        device=self._device,
        dtype=self._dtype,
    )

    iterator = train_loader
    if progress_bar:
        from tqdm import tqdm

        iterator = tqdm(train_loader, desc="Fitting ELLA")

    for step, batch in enumerate(iterator, start=1):
        x, y = self._batch(batch)
        features, output = self._features(x)
        if self.likelihood == Likelihood.REGRESSION:
            regression_targets(y, output)
            self.feature_gram += torch.einsum(
                "bck,bcl->kl", features, features
            ).detach()
        else:
            probs = output.softmax(dim=-1)
            hessian = torch.diag_embed(probs) - probs.unsqueeze(
                -1
            ) * probs.unsqueeze(-2)
            self.feature_gram += torch.einsum(
                "bck,bcd,bdl->kl", features, hessian, features
            ).detach()
        self.n_data += output.shape[0]
        self.fit_history_["processed_examples"].append(float(self.n_data))
        if val_loader is not None and val_steps and step % val_steps == 0:
            self._build_precision()
            self._fitted = True
            self.fit_history_["val_nll"].append(
                float(self._validation_nll(val_loader))
            )

    if self.n_data == 0:
        raise ValueError("ELLA fit requires a nonempty training loader.")

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

_feature_variance

_feature_variance(features: Tensor) -> Tensor

Return per-input output covariance from projected features.

Parameters:

• features (Tensor) –

  Projected Jacobians shaped [batch, outputs, n_eigenvalues].

Returns:

• Tensor –

  Latent covariance shaped [batch, outputs, outputs] without observation noise, using the current posterior precision factor.

Source code in laplace/baselaplace.py

def _feature_variance(self, features: torch.Tensor) -> torch.Tensor:
    """Return per-input output covariance from projected features.

    Parameters
    ----------
    features : torch.Tensor
        Projected Jacobians shaped ``[batch, outputs, n_eigenvalues]``.

    Returns
    -------
    torch.Tensor
        Latent covariance shaped ``[batch, outputs, outputs]`` without
        observation noise, using the current posterior precision factor.
    """
    solved = torch.cholesky_solve(features.transpose(1, 2), self.chol_precision)
    return features @ solved

_feature_covariance

_feature_covariance(features: Tensor) -> Tensor

Return joint output covariance from projected features.

Parameters:

• features (Tensor) –

  Projected Jacobians shaped [batch, outputs, n_eigenvalues].

Returns:

• Tensor –

  Latent covariance shaped [batch * outputs, batch * outputs] without observation noise. Coordinates are ordered by input, then output.

Source code in laplace/baselaplace.py

def _feature_covariance(self, features: torch.Tensor) -> torch.Tensor:
    """Return joint output covariance from projected features.

    Parameters
    ----------
    features : torch.Tensor
        Projected Jacobians shaped ``[batch, outputs, n_eigenvalues]``.

    Returns
    -------
    torch.Tensor
        Latent covariance shaped ``[batch * outputs, batch * outputs]``
        without observation noise. Coordinates are ordered by input,
        then output.
    """
    flat = features.reshape(-1, self.n_eigenvalues)
    solved = torch.cholesky_solve(flat.T, self.chol_precision)
    return flat @ solved

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

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)
    return self._feature_variance(jacobians @ self.dual_directions)

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

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)
    return self._feature_covariance(jacobians @ self.dual_directions)

_glm_predictive_distribution

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

Compute model outputs and ELLA 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 ELLA 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.
    """
    x = to_device(x, self._device, self._dtype)
    features, mean = self._features(x)
    covariance = (
        self._feature_covariance(features)
        if joint
        else self._feature_variance(features)
    )
    if joint:
        mean = mean.flatten()
    if not self.enable_backprop:
        return mean.detach(), covariance.detach()
    return mean, covariance

_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

optimize_prior_precision

optimize_prior_precision(
    pred_type: PredType | str = PredType.GP,
    method: TuningMethod | str = TuningMethod.GRIDSEARCH,
    n_steps: int = 100,
    lr: float = 0.1,
    init_prior_prec: float | Tensor = 1.0,
    prior_structure: (
        PriorStructure | str
    ) = PriorStructure.SCALAR,
    val_loader: DataLoader | None = None,
    loss: Any = None,
    log_prior_prec_min: float = -4,
    log_prior_prec_max: float = 4,
    grid_size: int = 100,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    verbose: bool = False,
    progress_bar: bool = False,
) -> None

Choose scalar prior precision by validation grid search.

Parameters:

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

  Must be "gp".

• method (TuningMethod or str, default: TuningMethod.GRIDSEARCH ) –

  Must be "gridsearch".

• n_steps (int, default: 100 ) –

  Compatibility argument; only the default is supported and it does not determine grid size.

• lr (float, default: 0.1 ) –

  Compatibility argument; only the default is supported.

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

  Compatibility argument; must be the scalar default. Grid search evaluates the specified precision range instead.

• prior_structure (PriorStructure or str, default: PriorStructure.SCALAR ) –

  Only scalar prior precision is supported.

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

  Validation data; required for grid search.

• loss (Any, default: None ) –

  Must be None; candidates are scored by predictive NLL.

• log_prior_prec_min (float, default: -4 ) –

  Base-10 logarithm of the smallest grid precision.

• log_prior_prec_max (float, default: 4 ) –

  Base-10 logarithm of the largest grid precision.

• grid_size (int, default: 100 ) –

  Number of candidate precisions.

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

  Classification validation link; only probit is supported.

• n_samples (int, default: 100 ) –

  Compatibility argument; only the default is supported. Validation scoring does not use Monte Carlo samples.

• verbose (bool, default: False ) –

  Compatibility argument; must be false.

• progress_bar (bool, default: False ) –

  Compatibility argument; must be false.

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the prediction type is not "gp", validation data are missing, or the grid size or candidate precisions are invalid.

• NotImplementedError –

  If marginal-likelihood tuning or nondefault compatibility arguments are requested.

Notes

Other arguments are accepted for interface compatibility but only their defaults are supported by ELLA.

Source code in laplace/baselaplace.py

def optimize_prior_precision(
    self,
    pred_type: PredType | str = PredType.GP,
    method: TuningMethod | str = TuningMethod.GRIDSEARCH,
    n_steps: int = 100,
    lr: float = 1e-1,
    init_prior_prec: float | torch.Tensor = 1.0,
    prior_structure: PriorStructure | str = PriorStructure.SCALAR,
    val_loader: DataLoader | None = None,
    loss: Any = None,
    log_prior_prec_min: float = -4,
    log_prior_prec_max: float = 4,
    grid_size: int = 100,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    verbose: bool = False,
    progress_bar: bool = False,
) -> None:
    """Choose scalar prior precision by validation grid search.

    Parameters
    ----------
    pred_type : PredType or str, default=PredType.GP
        Must be ``"gp"``.
    method : TuningMethod or str, default=TuningMethod.GRIDSEARCH
        Must be ``"gridsearch"``.
    n_steps : int, default=100
        Compatibility argument; only the default is supported and it does
        not determine grid size.
    lr : float, default=0.1
        Compatibility argument; only the default is supported.
    init_prior_prec : float or torch.Tensor, default=1.0
        Compatibility argument; must be the scalar default. Grid search
        evaluates the specified precision range instead.
    prior_structure : PriorStructure or str, default=PriorStructure.SCALAR
        Only scalar prior precision is supported.
    val_loader : DataLoader or None, default=None
        Validation data; required for grid search.
    loss : Any, default=None
        Must be ``None``; candidates are scored by predictive NLL.
    log_prior_prec_min : float, default=-4
        Base-10 logarithm of the smallest grid precision.
    log_prior_prec_max : float, default=4
        Base-10 logarithm of the largest grid precision.
    grid_size : int, default=100
        Number of candidate precisions.
    link_approx : LinkApprox or str, default=LinkApprox.PROBIT
        Classification validation link; only probit is supported.
    n_samples : int, default=100
        Compatibility argument; only the default is supported. Validation
        scoring does not use Monte Carlo samples.
    verbose : bool, default=False
        Compatibility argument; must be false.
    progress_bar : bool, default=False
        Compatibility argument; must be false.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the prediction type is not ``"gp"``, validation data are
        missing, or the grid size or candidate precisions are invalid.
    NotImplementedError
        If marginal-likelihood tuning or nondefault compatibility
        arguments are requested.

    Notes
    -----
    Other arguments are accepted for interface compatibility but only
    their defaults are supported by ELLA.
    """
    self._check_fitted()
    if pred_type != PredType.GP:
        raise ValueError("ELLA supports only pred_type='gp'.")
    if (
        n_steps != 100
        or lr != 1e-1
        or torch.as_tensor(init_prior_prec).numel() != 1
        or torch.as_tensor(init_prior_prec).item() != 1.0
        or prior_structure != PriorStructure.SCALAR
        or loss is not None
        or link_approx != LinkApprox.PROBIT
        or n_samples != 100
        or verbose
        or progress_bar
    ):
        raise NotImplementedError(
            "ELLA supports validation grid search without marginal-likelihood options."
        )
    if method != TuningMethod.GRIDSEARCH:
        raise NotImplementedError(
            "ELLA supports validation grid search, not marginal likelihood."
        )
    if val_loader is None:
        raise ValueError("gridsearch requires val_loader.")
    if (
        not isinstance(grid_size, Integral)
        or isinstance(grid_size, bool)
        or grid_size <= 0
    ):
        raise ValueError("grid_size must be positive integer.")
    grid = torch.logspace(log_prior_prec_min, log_prior_prec_max, grid_size)
    self.optimize_hyperparameters(val_loader, grid)

optimize_hyperparameters

optimize_hyperparameters(
    val_loader: DataLoader, grid
) -> None

Select the prior (and optionally noise) with lowest validation NLL.

Parameters:

• val_loader (DataLoader) –

  Validation data used to score every candidate.

• grid (iterable) –

  Candidate prior precisions or (prior_precision, sigma_noise) pairs for regression.

Raises:

• ValueError –

  If the grid is empty or all candidate scores are nonfinite.

Source code in laplace/baselaplace.py

def optimize_hyperparameters(self, val_loader: DataLoader, grid) -> None:
    """Select the prior (and optionally noise) with lowest validation NLL.

    Parameters
    ----------
    val_loader : DataLoader
        Validation data used to score every candidate.
    grid : iterable
        Candidate prior precisions or ``(prior_precision, sigma_noise)``
        pairs for regression.

    Raises
    ------
    ValueError
        If the grid is empty or all candidate scores are nonfinite.
    """
    self._check_fitted()
    initial_prior = self.prior_precision.detach().clone()
    initial_noise = self.sigma_noise.detach().clone()
    best_score = float("inf")
    best = None
    has_candidates = False
    for candidate in grid:
        has_candidates = True
        if isinstance(candidate, (tuple, list)):
            self.prior_precision, self.sigma_noise = candidate
        else:
            self.prior_precision = candidate
        score = float(self._validation_nll(val_loader))
        self.fit_history_["tuning"].append(
            {
                "prior_precision": float(self.prior_precision),
                "sigma_noise": float(self.sigma_noise),
                "val_nll": score,
            }
        )
        if isfinite(score) and score < best_score:
            best_score, best = score, candidate
    if best is None:
        self.prior_precision, self.sigma_noise = initial_prior, initial_noise
        self._refresh_posterior()
        if not has_candidates:
            raise ValueError("Hyperparameter grid is empty.")
        raise ValueError("No finite validation NLL was found in the grid.")
    if isinstance(best, (tuple, list)):
        self.prior_precision, self.sigma_noise = best
    else:
        self.prior_precision = best
    self._refresh_posterior()

log_marginal_likelihood

log_marginal_likelihood(*args, **kwargs)

Reject marginal-likelihood evaluation for the ELLA 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. ELLA does not define the subset-of-data GP marginal likelihood; use validation-based hyperparameter tuning.

Source code in laplace/baselaplace.py

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

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

    Raises
    ------
    NotImplementedError
        Always raised. ELLA does not define the subset-of-data GP
        marginal likelihood; use validation-based hyperparameter tuning.
    """
    raise NotImplementedError(
        "ELLA does not use FunctionalLaplace's subset-of-data marginal likelihood."
    )

state_dict

state_dict() -> dict[str, Any]

Return a checkpoint containing the fitted ELLA posterior.

Returns:

• dict[str, Any] –

  Basis directions, projected curvature, hyperparameters, fitting history, seed, model fingerprint, and backend/input settings. The posterior precision factor is rebuilt when loading.

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 the fitted ELLA posterior.

    Returns
    -------
    dict[str, Any]
        Basis directions, projected curvature, hyperparameters, fitting
        history, seed, model fingerprint, and backend/input settings.
        The posterior precision factor is rebuilt when loading.

    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,
        "subsample_size": self.subsample_size,
        "n_eigenvalues": self.n_eigenvalues,
        "dual_directions": self.dual_directions.detach().clone()
        if self.dual_directions is not None
        else None,
        "feature_gram": self.feature_gram.detach().clone()
        if self.feature_gram is not None
        else None,
        "prior_precision": self.prior_precision,
        "sigma_noise": self.sigma_noise,
        "temperature": self.temperature,
        "n_data": self.n_data,
        "n_outputs": self.n_outputs,
        "fitted": self._fitted,
        "fit_history": deepcopy(self.fit_history_),
        "seed": self.seed,
        **self._checkpoint_configuration(),
    }

load_state_dict

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

Restore a fitted ELLA 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 ELLA 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("ELLA 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 ELLA 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.subsample_size = state_dict["subsample_size"]
    self.n_eigenvalues = state_dict["n_eigenvalues"]
    self.dual_directions = state_dict["dual_directions"].to(self._device)
    self.feature_gram = state_dict["feature_gram"].to(self._device)
    self.prior_precision = state_dict["prior_precision"]
    self.sigma_noise = state_dict["sigma_noise"]
    self.temperature = state_dict["temperature"]
    self.n_data = state_dict["n_data"]
    self.n_outputs = state_dict["n_outputs"]
    setattr(self.model, "output_size", self.n_outputs)
    self._fitted = state_dict["fitted"]
    self.fit_history_ = deepcopy(state_dict["fit_history"])
    self._build_precision()
    self.model.eval()