loss_fns #

Loss functions.

Classes:

Name	Description
`BCELoss`
`BCEWithLogitsLoss`
`CrossEntropyLoss`
`L1Loss`
`MSELoss`
`NLLLoss`
`VariationalFreeEnergy`	Variational Free Energy Loss.

Functions:

Name	Description
`inputs_and_expanded_targets`	Ensure loss can be computed with additional dimensions of (sampled) predictions in inputs.

Attributes:

Name	Type	Description
`NegativeELBO`

NegativeELBO #

NegativeELBO = VariationalFreeEnergy

BCELoss #

Bases: BCELoss

Methods:

Name	Description
`forward`

forward #

forward(input: Tensor, target: Tensor) -> Tensor

BCEWithLogitsLoss #

Bases: BCEWithLogitsLoss

Methods:

Name	Description
`forward`

forward #

forward(input: Tensor, target: Tensor) -> Tensor

CrossEntropyLoss #

Bases: CrossEntropyLoss

Methods:

Name	Description
`forward`

forward #

forward(input: Tensor, target: Tensor) -> Tensor

L1Loss #

Bases: L1Loss

Methods:

Name	Description
`forward`

forward #

forward(input: Tensor, target: Tensor) -> Tensor

MSELoss #

Bases: MSELoss

Methods:

Name	Description
`forward`

forward #

forward(input: Tensor, target: Tensor) -> Tensor

NLLLoss #

Bases: NLLLoss

Methods:

Name	Description
`forward`

forward #

forward(input: Tensor, target: Tensor) -> Tensor

VariationalFreeEnergy #

VariationalFreeEnergy(
    nll: _Loss,
    model: BNNMixin,
    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`	`BNNMixin`	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, ""]

inputs_and_expanded_targets #

inputs_and_expanded_targets(inputs, targets)

Ensure loss can be computed with additional dimensions of (sampled) predictions in inputs.

Parameters:

Name	Type	Description	Default
`inputs`		Inputs (predictions).	required
`targets`		Targets.	required