torch_blue.vi.NormalNormalDivergence

class torch_blue.vi.NormalNormalDivergence

Bases: KullbackLeiblerModule

Kullback-Leibler divergence between two normal distributions.

Calculates the KL-Divergence between a MeanFieldNormalPrior and a MeanFieldNormalVarDist.

static forward(prior_mean: torch.Tensor | float, prior_log_std: torch.Tensor | float, variational_mean: torch.Tensor, variational_log_std: torch.Tensor) → torch.Tensor

Calculate the Kullback-Leibler divergence.

All input tensors must have the same shape.

This is not affected by _globals._USE_NORM_CONSTANTS.

Parameters:

  • prior_mean (Union[Tensor, float]) – Means of the prior distribution.

  • prior_log_std (Union[Tensor, float]) – Log standard deviations of the prior distribution.

  • variational_mean (Tensor) – Means of the variational distribution.

  • variational_log_std (Tensor) – Standard deviations of the variational distribution.

Returns:

The KL-Divergence of the two distributions.

Return type:

Tensor