Skip to content

8. Neural Network

8.1 conv1d

This module contains the code for the bayesian Conv1D.

8.1.1 Conv1D(input_channels, output_channels, kernel_size, stride=1, padding='VALID', dilation=1, groups=1, data_format='NWC', weights_distribution=None, bias_distribution=None, **kwargs)

This class is the bayesian implementation of the Conv1D class.

Initializes a Bayesian Conv1D layer.

Parameters:

Name Type Description Default
input_channels int

The number of channels in the input image.

 required
output_channels int

The number of channels produced by the convolution.

 required
kernel_size int

The size of the convolving kernel.

 required
stride Union[int, list[int]]

The stride of the convolution.

 1
padding str

The padding added to all four sides of the input. Can be 'VALID' or 'SAME'.

 'VALID'
dilation Union[int, list[int]]

The spacing between kernel elements.

 1
groups int

The number of blocked connections from input channels to output channels.

 1
data_format Optional[str]

The data format for the convolution, either 'NWC' or 'NCW'.

 'NWC'
weights_distribution Optional[GaussianDistribution]

The Gaussian distribution for the weights, if applicable.

 None
bias_distribution Optional[GaussianDistribution]

The Gaussian distribution for the bias, if applicable.

 None
**kwargs Any

Additional keyword arguments.

 {}
Source code in illia/nn/tf/conv1d.py 
23
24
25
26
27
28
29
30
31
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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def __init__(
    self,
    input_channels: int,
    output_channels: int,
    kernel_size: int,
    stride: Union[int, list[int]] = 1,
    padding: str = "VALID",
    dilation: Union[int, list[int]] = 1,
    groups: int = 1,
    data_format: Optional[str] = "NWC",
    weights_distribution: Optional[GaussianDistribution] = None,
    bias_distribution: Optional[GaussianDistribution] = None,
    **kwargs: Any,
) -> None:
    """
    Initializes a Bayesian Conv1D layer.

    Args:
        input_channels: The number of channels in the input image.
        output_channels: The number of channels produced by the
            convolution.
        kernel_size: The size of the convolving kernel.
        stride: The stride of the convolution.
        padding: The padding added to all four sides of the input.
            Can be 'VALID' or 'SAME'.
        dilation: The spacing between kernel elements.
        groups: The number of blocked connections from input
            channels to output channels.
        data_format: The data format for the convolution, either
            'NWC' or 'NCW'.
        weights_distribution: The Gaussian distribution for the
            weights, if applicable.
        bias_distribution: The Gaussian distribution for the bias,
            if applicable.
        **kwargs: Additional keyword arguments.
    """

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

    # Check data format
    self._check_params(kernel_size, groups, stride, dilation, data_format)

    # Set attributes
    self.input_channels = input_channels
    self.output_channels = output_channels
    self.kernel_size = kernel_size
    self.stride = stride
    self.padding = padding
    self.dilation = dilation
    self.groups = groups

    # Adjust the weights distribution based on the channel format
    self.data_format = (
        "NWC" if data_format is None or data_format == "NWC" else "NCW"
    )

    # Get the weights distribution shape, needs to be channel last
    self._weights_distribution_shape = (
        input_channels // groups,
        kernel_size,
        output_channels,
    )

    # Set weights distribution
    if weights_distribution is None:
        self.weights_distribution = GaussianDistribution(
            self._weights_distribution_shape
        )
    else:
        self.weights_distribution = weights_distribution

    # Set bias distribution
    if bias_distribution is None:
        self.bias_distribution = GaussianDistribution((output_channels,))
    else:
        self.bias_distribution = bias_distribution

8.1.1.1 call(inputs)

Performs a forward pass through the Bayesian Convolution 1D layer. If the layer is not frozen, it samples weights and bias from their respective distributions. If the layer is frozen and the weights or bias are not initialized, it also performs sampling.

Parameters:

Name Type Description Default
inputs Tensor

Input tensor to the layer.

 required

Returns:

Type Description
Tensor

Output tensor after passing through the layer.
Source code in illia/nn/tf/conv1d.py 
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
def call(self, inputs: tf.Tensor) -> tf.Tensor:
    """
    Performs a forward pass through the Bayesian Convolution 1D
    layer. If the layer is not frozen, it samples weights and bias
    from their respective distributions. If the layer is frozen
    and the weights or bias are not initialized, it also performs
    sampling.

    Args:
        inputs: Input tensor to the layer.

    Returns:
        Output tensor after passing through the layer.
    """

    # Check if layer is frozen
    if not self.frozen:
        self.w = self.weights_distribution.sample()
        self.b = self.bias_distribution.sample()
    elif self.w is None or self.b is None:
        raise ValueError(
            "Module has been frozen with undefined weights and/or bias."
        )

    # Compute outputs
    outputs: tf.Tensor = self._conv1d(
        inputs=inputs,
        weight=self.w,
        stride=self.stride,
        padding=self.padding,
        data_format=self.data_format,
        dilation=self.dilation,
    )

    # Add bias
    outputs = tf.nn.bias_add(
        value=outputs,
        bias=self.b,
        data_format="N..C" if self.data_format == "NWC" else "NC..",
    )

    return outputs

8.1.1.2 freeze()

Freezes the current module and all submodules that are instances of BayesianModule. Sets the frozen state to True.

Source code in illia/nn/tf/conv1d.py 
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
def freeze(self) -> None:
    """
    Freezes the current module and all submodules that are instances
    of BayesianModule. Sets the frozen state to True.
    """

    # Set indicator
    self.frozen = True

    # Sample weights if they are undefined
    if self.w is None:
        self.w = self.weights_distribution.sample()

    # Sample bias is they are undefined
    if self.b is None:
        self.b = self.bias_distribution.sample()

8.1.1.3 kl_cost()

Computes the Kullback-Leibler (KL) divergence cost for the layer's weights and bias.

Returns:

Type Description
Tensor

Tuple containing KL divergence cost and total number of
int

parameters.
Source code in illia/nn/tf/conv1d.py 
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
def kl_cost(self) -> tuple[tf.Tensor, int]:
    """
    Computes the Kullback-Leibler (KL) divergence cost for the
    layer's weights and bias.

    Returns:
        Tuple containing KL divergence cost and total number of
        parameters.
    """

    # Compute log probs
    log_probs: tf.Tensor = self.weights_distribution.log_prob(
        self.w
    ) + self.bias_distribution.log_prob(self.b)

    # Compute the number of parameters
    num_params: int = (
        self.weights_distribution.num_params + self.bias_distribution.num_params
    )

    return log_probs, num_params

8.2 conv2d

This module contains the code for the bayesian Conv2D.

8.2.1 Conv2D(input_channels, output_channels, kernel_size, stride=1, padding='VALID', dilation=None, groups=1, data_format='NHWC', weights_distribution=None, bias_distribution=None, **kwargs)

This class is the bayesian implementation of the Conv2D class.

Initializes a Bayesian Conv2D layer.

Parameters:

Name Type Description Default
input_channels int

The number of channels in the input image.

 required
output_channels int

The number of channels produced by the convolution.

 required
kernel_size Union[int, list[int]]

The size of the convolving kernel.

 required
stride Union[int, list[int]]

The stride of the convolution.

 1
padding Union[str, list[int]]

The padding added to all four sides of the input. Can be 'VALID' or 'SAME'.

 'VALID'
dilation Optional[Union[int, list[int]]]

The spacing between kernel elements.

 None
groups int

The number of blocked connections from input channels to output channels.

 1
data_format Optional[str]

The data format for the convolution, either 'NHWC' or 'NCHW'.

 'NHWC'
weights_distribution Optional[GaussianDistribution]

The Gaussian distribution for the weights, if applicable.

 None
bias_distribution Optional[GaussianDistribution]

The Gaussian distribution for the bias, if applicable.

 None
**kwargs Any

Additional keyword arguments.

 {}
Source code in illia/nn/tf/conv2d.py 
 23
 24
 25
 26
 27
 28
 29
 30
 31
 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
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def __init__(
    self,
    input_channels: int,
    output_channels: int,
    kernel_size: Union[int, list[int]],
    stride: Union[int, list[int]] = 1,
    padding: Union[str, list[int]] = "VALID",
    dilation: Optional[Union[int, list[int]]] = None,
    groups: int = 1,
    data_format: Optional[str] = "NHWC",
    weights_distribution: Optional[GaussianDistribution] = None,
    bias_distribution: Optional[GaussianDistribution] = None,
    **kwargs: Any,
) -> None:
    """
    Initializes a Bayesian Conv2D layer.

    Args:
        input_channels: The number of channels in the input image.
        output_channels: The number of channels produced by the
            convolution.
        kernel_size: The size of the convolving kernel.
        stride: The stride of the convolution.
        padding: The padding added to all four sides of the input.
            Can be 'VALID' or 'SAME'.
        dilation: The spacing between kernel elements.
        groups: The number of blocked connections from input channels
            to output channels.
        data_format: The data format for the convolution, either
            'NHWC' or 'NCHW'.
        weights_distribution: The Gaussian distribution for the
            weights, if applicable.
        bias_distribution: The Gaussian distribution for the bias,
            if applicable.
        **kwargs: Additional keyword arguments.
    """

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

    # Check data format
    self._check_params(kernel_size, groups, stride, dilation, data_format)

    # Set attributes
    self.input_channels = input_channels
    self.output_channels = output_channels
    self.kernel_size = kernel_size
    self.stride = stride
    self.padding = padding
    self.dilation = dilation
    self.groups = groups

    # Check if kernel_size is a list and unpack it if necessary
    kernel_shape = (
        kernel_size if isinstance(kernel_size, list) else [kernel_size, kernel_size]
    )

    # Adjust the weights distribution based on the channel format
    self.data_format = (
        "NHWC" if data_format is None or data_format == "NHWC" else "NCHW"
    )

    # Set the weights distribution shape
    self._weights_distribution_shape = (
        input_channels // groups,
        *kernel_shape,
        output_channels,
    )

    # Set weights distribution
    if weights_distribution is None:
        self.weights_distribution = GaussianDistribution(
            shape=self._weights_distribution_shape
        )
    else:
        self.weights_distribution = weights_distribution

    # Set bias distribution
    if bias_distribution is None:
        self.bias_distribution = GaussianDistribution((output_channels,))
    else:
        self.bias_distribution = bias_distribution

8.2.1.1 call(inputs)

Performs a forward pass through the Bayesian Convolution 2D layer. If the layer is not frozen, it samples weights and bias from their respective distributions. If the layer is frozen and the weights or bias are not initialized, it also performs sampling.

Parameters:

Name Type Description Default
inputs Tensor

Input tensor to the layer. Dimensions: [batch, input channels, input width, input height].

 required

Returns:

Type Description
Tensor

Output tensor after passing through the layer. Dimensions: [batch, output channels, output width, output height].
Source code in illia/nn/tf/conv2d.py 
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
def call(self, inputs: tf.Tensor) -> tf.Tensor:
    """
    Performs a forward pass through the Bayesian Convolution 2D
    layer. If the layer is not frozen, it samples weights and bias
    from their respective distributions. If the layer is frozen
    and the weights or bias are not initialized, it also performs
    sampling.

    Args:
        inputs: Input tensor to the layer. Dimensions: [batch,
            input channels, input width, input height].

    Returns:
        Output tensor after passing through the layer. Dimensions:
            [batch, output channels, output width, output height].
    """

    # Check if layer is frozen
    if not self.frozen:
        self.w = self.weights_distribution.sample()
        self.b = self.bias_distribution.sample()
    elif self.w is None or self.b is None:
        raise ValueError(
            "Module has been frozen with undefined weights and/or bias."
        )

    # Compute outputs
    outputs: tf.Tensor = self._conv2d(
        inputs=inputs,
        weight=self.w,
        stride=self.stride,
        padding=self.padding,
        data_format=self.data_format,
        dilation=self.dilation,
    )

    # Add bias
    outputs = tf.nn.bias_add(
        value=outputs,
        bias=self.b,
        data_format="N..C" if self.data_format == "NHWC" else "NC..",
    )

    return outputs

8.2.1.2 freeze()

Freezes the current module and all submodules that are instances of BayesianModule. Sets the frozen state to True.

Source code in illia/nn/tf/conv2d.py 
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
def freeze(self) -> None:
    """
    Freezes the current module and all submodules that are instances
    of BayesianModule. Sets the frozen state to True.
    """

    # Set indicator
    self.frozen = True

    # Sample weights if they are undefined
    if self.w is None:
        self.w = self.weights_distribution.sample()

    # Sample bias is they are undefined
    if self.b is None:
        self.b = self.bias_distribution.sample()

8.2.1.3 kl_cost()

Computes the Kullback-Leibler (KL) divergence cost for the layer's weights and bias.

Returns:

Type Description
Tensor

Tuple containing KL divergence cost and total number of
int

parameters.
Source code in illia/nn/tf/conv2d.py 
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
def kl_cost(self) -> tuple[tf.Tensor, int]:
    """
    Computes the Kullback-Leibler (KL) divergence cost for the
    layer's weights and bias.

    Returns:
        Tuple containing KL divergence cost and total number of
        parameters.
    """

    # Compute log probs
    log_probs: tf.Tensor = self.weights_distribution.log_prob(
        self.w
    ) + self.bias_distribution.log_prob(self.b)

    # Compute the number of parameters
    num_params: int = (
        self.weights_distribution.num_params + self.bias_distribution.num_params
    )

    return log_probs, num_params

8.3 embedding

This module contains the code for Embedding Bayesian layer.

8.3.1 Embedding(num_embeddings, embeddings_dim, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False, weights_distribution=None, **kwargs)

This class is the bayesian implementation of the Embedding class.

This method is the constructor of the embedding class.

Parameters:

Name Type Description Default
num_embeddings int

Size of the dictionary of embeddings.

 required
embeddings_dim int

The size of each embedding vector.

 required
padding_idx Optional[int]

If specified, the entries at padding_idx do not contribute to the gradient.

 None
max_norm Optional[float]

If given, each embedding vector with norm larger than max_norm is renormalized to have norm max_norm.

 None
norm_type float

The p of the p-norm to compute for the max_norm option.

 2.0
scale_grad_by_freq bool

If given, this will scale gradients by the inverse of frequency of the words in the mini-batch.

 False
sparse bool

If True, gradient w.r.t. weight matrix will be a sparse tensor.

 False
weights_distribution Optional[GaussianDistribution]

The Gaussian distribution for the weights, if applicable.

 None
**kwargs Any

Additional keyword arguments.

 {}
Source code in illia/nn/tf/embedding.py 
23
24
25
26
27
28
29
30
31
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
74
75
def __init__(
    self,
    num_embeddings: int,
    embeddings_dim: int,
    padding_idx: Optional[int] = None,
    max_norm: Optional[float] = None,
    norm_type: float = 2.0,
    scale_grad_by_freq: bool = False,
    sparse: bool = False,
    weights_distribution: Optional[GaussianDistribution] = None,
    **kwargs: Any,
) -> None:
    """
    This method is the constructor of the embedding class.

    Args:
        num_embeddings: Size of the dictionary of embeddings.
        embeddings_dim: The size of each embedding vector.
        padding_idx: If specified, the entries at padding_idx do
            not contribute to the gradient.
        max_norm: If given, each embedding vector with norm larger
            than max_norm is renormalized to have norm max_norm.
        norm_type: The p of the p-norm to compute for the max_norm
            option.
        scale_grad_by_freq: If given, this will scale gradients by
            the inverse of frequency of the words in the
            mini-batch.
        sparse: If True, gradient w.r.t. weight matrix will be a
            sparse tensor.
        weights_distribution: The Gaussian distribution for the
            weights, if applicable.
        **kwargs: Additional keyword arguments.
    """

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

    # Set atributtes
    self.num_embeddings = num_embeddings
    self.embeddings_dim = embeddings_dim
    self.padding_idx = padding_idx
    self.max_norm = max_norm
    self.norm_type = norm_type
    self.scale_grad_by_freq = scale_grad_by_freq
    self.sparse = sparse

    # Set weights distribution
    if weights_distribution is None:
        self.weights_distribution = GaussianDistribution(
            (num_embeddings, embeddings_dim)
        )
    else:
        self.weights_distribution = weights_distribution

8.3.1.1 call(inputs)

Performs a forward pass through the Bayesian Embedding layer.

Samples weights from their posterior distributions if the layer is not frozen. If frozen and not initialized, samples them once.

Parameters:

Name Type Description Default
inputs Tensor

input tensor. Dimensions: [batch, *].

 required

Raises:

Type Description
ValueError

Module has been frozen with undefined weights.

Returns:

Type Description
Tensor

Output tensor after linear transformation.
Source code in illia/nn/tf/embedding.py 
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
def call(self, inputs: tf.Tensor) -> tf.Tensor:
    """
    Performs a forward pass through the Bayesian Embedding layer.

    Samples weights from their posterior distributions if
    the layer is not frozen. If frozen and not initialized, samples
    them once.

    Args:
        inputs: input tensor. Dimensions: [batch, *].

    Raises:
        ValueError: Module has been frozen with undefined weights.

    Returns:
        Output tensor after linear transformation.
    """

    # Check if layer is frozen
    if not self.frozen:
        self.w = self.weights_distribution.sample()
    elif self.w is None:
        raise ValueError("Module has been frozen with undefined weights.")

    # Compute outputs
    outputs: tf.Tensor = self._embedding(
        inputs,
        self.w,
        self.padding_idx,
        self.max_norm,
        self.norm_type,
        self.sparse,
    )

    return outputs

8.3.1.2 freeze()

Freezes the current module and all submodules that are instances of BayesianModule. Sets the frozen state to True.

Source code in illia/nn/tf/embedding.py 
168
169
170
171
172
173
174
175
176
177
178
179
def freeze(self) -> None:
    """
    Freezes the current module and all submodules that are instances
    of BayesianModule. Sets the frozen state to True.
    """

    # Set indicator
    self.frozen = True

    # Sample weights if they are undefined
    if self.w is None:
        self.w = self.weights_distribution.sample()

8.3.1.3 kl_cost()

Computes the Kullback-Leibler (KL) divergence cost for the layer's weights.

Returns:

Type Description
Tensor

Tuple containing KL divergence cost and total number of
int

parameters.
Source code in illia/nn/tf/embedding.py 
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
def kl_cost(self) -> tuple[tf.Tensor, int]:
    """
    Computes the Kullback-Leibler (KL) divergence cost for the
    layer's weights.

    Returns:
        Tuple containing KL divergence cost and total number of
        parameters.
    """

    # Get log probs
    log_probs: tf.Tensor = self.weights_distribution.log_prob(self.w)

    # Get number of parameters
    num_params: int = self.weights_distribution.num_params

    return log_probs, num_params

8.4 linear

This module contains the code for Linear Bayesian layer.

8.4.1 Linear(input_size, output_size, weights_distribution=None, bias_distribution=None, **kwargs)

This class is the bayesian implementation of the Linear class.

This is the constructor of the Linear class.

Parameters:

Name Type Description Default
input_size int

Input size of the linear layer.

 required
output_size int

Output size of the linear layer.

 required
weights_distribution Optional[GaussianDistribution]

The Gaussian distribution for the weights, if applicable.

 None
bias_distribution Optional[GaussianDistribution]

The Gaussian distribution for the bias, if applicable.

 None
**kwargs Any

Additional keyword arguments.

 {}
Source code in illia/nn/tf/linear.py 
23
24
25
26
27
28
29
30
31
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
def __init__(
    self,
    input_size: int,
    output_size: int,
    weights_distribution: Optional[GaussianDistribution] = None,
    bias_distribution: Optional[GaussianDistribution] = None,
    **kwargs: Any,
) -> None:
    """
    This is the constructor of the Linear class.

    Args:
        input_size: Input size of the linear layer.
        output_size: Output size of the linear layer.
        weights_distribution: The Gaussian distribution for the
            weights, if applicable.
        bias_distribution: The Gaussian distribution for the bias,
            if applicable.
        **kwargs: Additional keyword arguments.
    """

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

    # Set parameters
    self.input_size = input_size
    self.output_size = output_size

    # Set weights distribution
    if weights_distribution is None:
        self.weights_distribution = GaussianDistribution((output_size, input_size))
    else:
        self.weights_distribution = weights_distribution

    # Set bias distribution
    if bias_distribution is None:
        self.bias_distribution = GaussianDistribution((output_size,))
    else:
        self.bias_distribution = bias_distribution

8.4.1.1 call(inputs)

Performs a forward pass through the Bayesian Linear layer.

Samples weights and bias from their posterior distributions if the layer is not frozen. If frozen and not initialized, samples them once.

Parameters:

Name Type Description Default
inputs Tensor

input tensor. Dimensions: [batch, *].

 required

Raises:

Type Description
ValueError

Module has been frozen with undefined weights.

Returns:

Type Description
Tensor

Output tensor after linear transformation.
Source code in illia/nn/tf/linear.py 
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
def call(self, inputs: tf.Tensor) -> tf.Tensor:
    """
    Performs a forward pass through the Bayesian Linear layer.

    Samples weights and bias from their posterior distributions if
    the layer is not frozen. If frozen and not initialized, samples
    them once.

    Args:
        inputs: input tensor. Dimensions: [batch, *].

    Raises:
        ValueError: Module has been frozen with undefined weights.

    Returns:
        Output tensor after linear transformation.
    """

    # Check if layer is frozen
    if not self.frozen:
        self.w = self.weights_distribution.sample()
        self.b = self.bias_distribution.sample()
    elif self.w is None or self.b is None:
        raise ValueError(
            "Module has been frozen with undefined weights and/or bias."
        )

    # Compute outputs
    outputs: tf.Tensor = tf.nn.bias_add(
        tf.linalg.matmul(inputs, self.w, transpose_b=True), self.b
    )

    return outputs

8.4.1.2 freeze()

Freezes the current module and all submodules that are instances of BayesianModule. Sets the frozen state to True.

Source code in illia/nn/tf/linear.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def freeze(self) -> None:
    """
    Freezes the current module and all submodules that are instances
    of BayesianModule. Sets the frozen state to True.
    """

    # Set indicator
    self.frozen = True

    # Sample weights if they are undefined
    if self.w is None:
        self.w = self.weights_distribution.sample()

    # Sample bias is they are undefined
    if self.b is None:
        self.b = self.bias_distribution.sample()

8.4.1.3 kl_cost()

Computes the Kullback-Leibler (KL) divergence cost for the layer's weights and bias.

Returns:

Type Description
Tensor

Tuple containing KL divergence cost and total number of
int

parameters.
Source code in illia/nn/tf/linear.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def kl_cost(self) -> tuple[tf.Tensor, int]:
    """
    Computes the Kullback-Leibler (KL) divergence cost for the
    layer's weights and bias.

    Returns:
        Tuple containing KL divergence cost and total number of
        parameters.
    """

    # Compute log probs
    log_probs: tf.Tensor = self.weights_distribution.log_prob(
        self.w
    ) + self.bias_distribution.log_prob(self.b)

    # Compute the number of parameters
    num_params: int = (
        self.weights_distribution.num_params + self.bias_distribution.num_params
    )

    return log_probs, num_params