loss_fns

Loss functions.

Classes:

Name	Description
VariationalFreeEnergy	Variational Free Energy Loss.

Attributes:

Name	Type	Description
NegativeELBO

NegativeELBO

NegativeELBO = VariationalFreeEnergy

VariationalFreeEnergy

VariationalFreeEnergy(
    nll: _Loss,
    model: BNNModule,
    prior_loc: Float[Tensor, "parameter"] | None = None,
    prior_scale: Float[Tensor, "parameter"] | None = None,
    kl_weight: float | None = 1.0,
    reduction: str = "mean",
)

Bases: Module

Variational Free Energy Loss.

Computes the variational free energy loss for variational inference with the Kullback-Leibler regularization term computed in weight space. This is also known as the negative evidence lower bound (ELBO).

Parameters:

Name	Type	Description	Default
nll	_Loss	Loss function defining the negative log-likelihood.	required
model	BNNModule	The probabilistic model.	required
prior_loc	Float[Tensor, 'parameter'] | None	Location(s) of the prior Gaussian distribution.	None
prior_scale	Float[Tensor, 'parameter'] | None	Scale(s) of the prior Gaussian distribution.	None
kl_weight	float | None	Weight for the KL divergence term. If None, chooses the weight inversely proportional to the number of mean parameters.	1.0
reduction	str	Specifies the reduction to apply to the output: ``'mean' | 'sum'. 'mean': the weighted mean of the output is taken, 'sum': the output will be summed.	'mean'

Methods:

Name	Description
forward

Attributes:

Name	Type	Description
kl_weight
model
nll
numel_mean_parameters
prior_loc
prior_scale
reduction

kl_weight

kl_weight = kl_weight

model

model = model

nll

nll = nll

numel_mean_parameters

numel_mean_parameters = sum(
    numel()
    for (name, param) in named_parameters()
    if requires_grad
    and "params." in name
    and "cov." not in name
)

prior_loc

prior_loc = prior_loc

prior_scale

prior_scale = prior_scale

reduction

reduction = reduction

forward

forward(
    input: Float[Tensor, "*sample batch in_feature"],
    target: Float[Tensor, "batch out_feature"],
) -> Float[Tensor, ""]