torch_blue.vi.distributions.Categorical

class torch_blue.vi.distributions.Categorical(input_type: str = 'logits')

Bases: torch_blue.vi.distributions.base.Distribution

Categorical distribution used as predictive distribution for classification tasks.

This distribution is implemented only as a predictive distribution. It interprets the model output either as logits (default) or as probabilities. In both cases, the distribution parameters are class probabilities with their sum normalized to one across classes. The last dimension of the model output is interpreted as the number of classes.

When used with KullbackLeiblerLoss or AnalyticalKullbackLeiblerLoss this distribution produces a loss related to cross-entropy loss.

Parameters:

input_type (str, default: "logits") – Whether to interpret the model output as logits (“logits”; i.e., log probabilities) or probabilities (“probs”).

Raises:

AssertionError – If input_type is neither “logits” nor “probs”.

predictive_parameters_from_samples(samples: torch.Tensor, eps: float = 1e-05) → torch.Tensor

Calculate predictive probabilities from samples.

Converts logits to probabilities if input_type is “logits” and normalizes along the class dimension.

Parameters:

  • samples (Tensor) – The model output as Tensor of shape (S, B, C), where S is the number of samples, B is the batch size, and C is the number of classes. The batch size dimension is optional.

  • eps (float, default: 1e-5) – Epsilon for numerical stability.

Returns:

The predictive class probabilities as Tensor of shape (B, C).

Return type:

Tensor

static log_prob_from_parameters(reference: torch.Tensor, parameters: torch.Tensor, eps: float = 1e-05) → torch.Tensor

Calculate the log probability of the label based on the class probabilities.

This is not affected by _globals._USE_NORM_CONSTANTS.

Parameters:

  • reference (Tensor) – The ground truth label as Tensor of shape (B,), where B is the batch size and may be one.

  • parameters (Tensor) – The predictive class probabilities as Tensor of shape (B, C) as returned by predictive_parameters_from_samples().

  • eps (float, default: 1e-5) – Epsilon for numerical stability.

Returns:

The log probability of the label under the predicted class probabilities. Shape: (1,).

Return type:

Tensor