Skip to content

4. Distributions

Base class for building distribution modules using PyTorch.

Provides a standardized interface for sampling, computing log probabilities, and reporting the number of parameters in custom probabilistic layers.

Notes

This class is abstract and should not be instantiated directly. Subclasses must implement all abstract methods to specify distribution behavior.

4.1 DistributionModule

Abstract base for PyTorch-based probabilistic distribution modules.

Defines a required interface for sampling, computing log-probabilities, and retrieving parameter counts. Subclasses must implement all abstract methods to provide specific distribution logic.

Notes

Avoid direct instantiation, this serves as a blueprint for derived classes.

4.1.1 log_prob(x=None) abstractmethod

Computes the log-probability of an input sample. If no sample is provided, a new one is drawn internally from the current distribution before computing the log-probability.

Parameters:

Name Type Description Default
x Optional[Tensor]

Optional sample tensor to evaluate.

 None

Returns:

Type Description
Tensor

Scalar tensor representing the computed log-probability.
Notes

This method supports both user-supplied samples and internally generated ones for convenience when evaluating likelihoods.

Source code in illia/distributions/torch/base.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
@abstractmethod
def log_prob(self, x: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Computes the log-probability of an input sample.
    If no sample is provided, a new one is drawn internally from the
    current distribution before computing the log-probability.

    Args:
        x: Optional sample tensor to evaluate.

    Returns:
        Scalar tensor representing the computed log-probability.

    Notes:
        This method supports both user-supplied samples and internally
        generated ones for convenience when evaluating likelihoods.
    """

4.1.2 num_params() abstractmethod

Returns the total number of learnable parameters in the distribution.

Returns:

Type Description
int

Integer representing the total number of learnable parameters.
Source code in illia/distributions/torch/base.py 
63
64
65
66
67
68
69
70
@abstractmethod
def num_params(self) -> int:
    """
    Returns the total number of learnable parameters in the distribution.

    Returns:
        Integer representing the total number of learnable parameters.
    """

4.1.3 sample() abstractmethod

Generates and returns a sample from the underlying distribution.

Returns:

Type Description
Tensor

Sample array matching the shape and structure defined by
Tensor

the distribution parameters.
Source code in illia/distributions/torch/base.py 
35
36
37
38
39
40
41
42
43
@abstractmethod
def sample(self) -> torch.Tensor:
    """
    Generates and returns a sample from the underlying distribution.

    Returns:
        Sample array matching the shape and structure defined by
        the distribution parameters.
    """

Defines a Gaussian (Normal) distribution using PyTorch with trainable mean and standard deviation parameters. Includes methods for sampling from the distribution and computing log-probabilities of given inputs.

4.2 GaussianDistribution(shape, mu_prior=0.0, std_prior=0.1, mu_init=0.0, rho_init=-7.0, **kwargs)

Learnable Gaussian distribution using PyTorch.

Represents a diagonal Gaussian distribution with trainable mean and standard deviation parameters. The standard deviation is derived from rho using a softplus transformation to ensure positivity.

Notes

Assumes a diagonal covariance matrix. KL divergence between distributions can be computed using log-probability differences obtained from log_prob.

Initializes the Gaussian distribution layer.

Parameters:

Name Type Description Default
shape tuple[int, ...]

Shape of the learnable parameters.

 required
mu_prior float

Mean of the Gaussian prior.

 0.0
std_prior float

Standard deviation of the prior.

 0.1
mu_init float

Initial value for the mean parameter.

 0.0
rho_init float

Initial value for the rho parameter.

 -7.0

Returns:

Type Description
None

None.
Source code in illia/distributions/torch/gaussian.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def __init__(
    self,
    shape: tuple[int, ...],
    mu_prior: float = 0.0,
    std_prior: float = 0.1,
    mu_init: float = 0.0,
    rho_init: float = -7.0,
    **kwargs: Any,
) -> None:
    """
    Initializes the Gaussian distribution layer.

    Args:
        shape: Shape of the learnable parameters.
        mu_prior: Mean of the Gaussian prior.
        std_prior: Standard deviation of the prior.
        mu_init: Initial value for the mean parameter.
        rho_init: Initial value for the rho parameter.

    Returns:
        None.
    """

    # Call super-class constructor
    super().__init__(**kwargs)

    # Set attributes
    self.shape = shape
    self.mu_init = mu_init
    self.rho_init = rho_init

    # Define priors
    self.register_buffer("mu_prior", torch.tensor([mu_prior]))
    self.register_buffer("std_prior", torch.tensor([std_prior]))

    # Define initial mu and rho
    self.mu: torch.Tensor = torch.nn.Parameter(
        torch.randn(self.shape).normal_(self.mu_init, 0.1)
    )
    self.rho: torch.Tensor = torch.nn.Parameter(
        torch.randn(self.shape).normal_(self.rho_init, 0.1)
    )

4.2.1 log_prob(x=None)

Computes the KL divergence between posterior and prior.

If no sample is given, one is drawn from the distribution.

Parameters:

Name Type Description Default
x Optional[Tensor]

Optional sample tensor. If None, generates a new sample.

 None

Returns:

Type Description
Tensor

A scalar tensor representing the KL divergence.
Source code in illia/distributions/torch/gaussian.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
@torch.jit.export
def log_prob(self, x: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Computes the KL divergence between posterior and prior.

    If no sample is given, one is drawn from the distribution.

    Args:
        x: Optional sample tensor. If None, generates a new sample.

    Returns:
        A scalar tensor representing the KL divergence.
    """

    # Sample if x is None
    if x is None:
        x = self.sample()

    # Define pi variable
    pi: torch.Tensor = torch.acos(torch.zeros(1)) * 2

    # Compute log priors
    log_prior = (
        -torch.log(torch.sqrt(2 * pi)).to(x.device)
        - torch.log(self.std_prior)
        - (((x - self.mu_prior) ** 2) / (2 * self.std_prior**2))
        - 0.5
    )

    # Compute sigma
    sigma: torch.Tensor = torch.log1p(torch.exp(self.rho)).to(x.device)

    # Compute log posteriors
    log_posteriors = (
        -torch.log(torch.sqrt(2 * pi)).to(x.device)
        - torch.log(sigma)
        - (((x - self.mu) ** 2) / (2 * sigma**2))
        - 0.5
    )

    # Compute final log probs
    log_probs = log_posteriors.sum() - log_prior.sum()

    return log_probs

4.2.2 num_params()

Returns the number of learnable parameters in the distribution.

Returns:

Type Description
int

Total number of parameters as an integer.
Source code in illia/distributions/torch/gaussian.py 
135
136
137
138
139
140
141
142
143
144
145
@torch.jit.export
@torch.no_grad()
def num_params(self) -> int:
    """
    Returns the number of learnable parameters in the distribution.

    Returns:
        Total number of parameters as an integer.
    """

    return len(self.mu.view(-1))

4.2.3 sample()

Draws a sample from the distribution using reparameterization.

Returns:

Type Description
Tensor

A sample tensor with the same shape as mu and rho.
Source code in illia/distributions/torch/gaussian.py 
75
76
77
78
79
80
81
82
83
84
85
86
87
88
@torch.jit.export
def sample(self) -> torch.Tensor:
    """
    Draws a sample from the distribution using reparameterization.

    Returns:
        A sample tensor with the same shape as `mu` and `rho`.
    """

    # Sampling with reparametrization trick
    eps: torch.Tensor = torch.randn_like(self.rho)
    sigma: torch.Tensor = torch.log1p(torch.exp(self.rho))

    return self.mu + sigma * eps