Skip to content

7. Distributions

Base class for distribution modules built using Tensorflow.

Provides a consistent interface for sampling, evaluating log probabilities, and retrieving the number of parameters in custom distribution layers.

Methods should be implemented by subclasses to define the actual sampling behavior, probability evaluation, and parameter reporting.

7.1 DistributionModule

Abstract base for Tensorflow-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.

7.1.1 num_params abstractmethod property

Returns the total number of learnable parameters in the distribution.

Returns:

Type Description
int

Integer representing the total number of learnable parameters.

7.1.2 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/tf/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[tf.Tensor] = None) -> tf.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.
    """

7.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/tf/base.py 
35
36
37
38
39
40
41
42
43
@abstractmethod
def sample(self) -> tf.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 Tensorflow with trainable mean and standard deviation parameters. Includes methods for sampling from the distribution and computing log-probabilities of given inputs.

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

Learnable Gaussian distribution using Tensorflow.

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

Prior mean for KL divergence computation.

 0.0
std_prior float

Prior std for KL divergence computation.

 0.1
mu_init float

Initial value for the mean.

 0.0
rho_init float

Initial value for the rho parameter.

 -7.0
**kwargs Any

Additional arguments passed to the base class.

 {}

Returns:

Type Description
None

None.
Source code in illia/distributions/tf/gaussian.py 
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
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: Prior mean for KL divergence computation.
        std_prior: Prior std for KL divergence computation.
        mu_init: Initial value for the mean.
        rho_init: Initial value for the rho parameter.
        **kwargs: Additional arguments passed to the base class.

    Returns:
        None.
    """

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

    # Set attributes
    self.shape = shape
    self.mu_prior_value = mu_prior
    self.std_prior_value = std_prior
    self.mu_init = mu_init
    self.rho_init = rho_init

    # Call build method
    self.build(self.shape)

7.2.1 num_params property

Returns the number of learnable parameters.

Returns:

Type Description
int

Total number of parameters in the distribution.

7.2.2 log_prob(x=None)

Computes KL divergence between posterior and prior.

If no sample is provided, one is drawn from the current distribution.

Parameters:

Name Type Description Default
x Optional[Tensor]

Optional input sample tensor.

 None

Returns:

Type Description
Tensor

A scalar tensor representing the KL divergence.
Source code in illia/distributions/tf/gaussian.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def log_prob(self, x: Optional[tf.Tensor] = None) -> tf.Tensor:
    """
    Computes KL divergence between posterior and prior.

    If no sample is provided, one is drawn from the current
    distribution.

    Args:
        x: Optional input sample tensor.

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

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

    # Define pi variable
    pi: tf.Tensor = tf.convert_to_tensor(math.pi)

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

    # Compute sigma
    sigma: tf.Tensor = tf.math.log1p(tf.math.exp(self.rho))

    # Compute log posteriors
    log_posteriors: tf.Tensor = (
        -tf.math.log(tf.math.sqrt(2 * pi))
        - tf.math.log(sigma)
        - (((x - self.mu) ** 2) / (2 * sigma**2))
        - 0.5
    )

    # Compute final log probs
    log_probs = tf.math.reduce_sum(log_posteriors) - tf.math.reduce_sum(log_prior)

    return log_probs

7.2.3 sample()

Draws a sample from the distribution using reparameterization.

Returns:

Type Description
Tensor

A sample tensor with the same shape as the parameters.
Source code in illia/distributions/tf/gaussian.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
def sample(self) -> tf.Tensor:
    """
    Draws a sample from the distribution using reparameterization.

    Returns:
        A sample tensor with the same shape as the parameters.
    """

    # Sampling with reparametrization trick
    eps: tf.Tensor = tf.random.normal(shape=self.rho.shape)
    sigma: tf.Tensor = tf.math.log1p(tf.math.exp(self.rho))

    return self.mu + sigma * eps