Skip to content

6. Neural Network

This module contains the code for the bayesian Conv1d.

6.1 Conv1d(input_channels, output_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, weights_distribution=None, bias_distribution=None, **kwargs)

This class is the bayesian implementation of the Conv1d class.

Definition of a Bayesian Convolution 1D layer.

Parameters:

Name Type Description Default
input_channels int

Number of channels in the input image.

 required
output_channels int

Number of channels produced by the convolution.

 required
kernel_size int

Size of the convolving kernel.

 required
stride int

Stride of the convolution. Deafults to 1.

 1
padding int

Padding added to all four sides of the input. Defaults to 0.

 0
dilation int

Spacing between kernel elements.

 1
groups int

Number of blocked connections from input channels to output channels. Defaults to 1.

 1
weights_distribution Optional[GaussianDistribution]

The distribution for the weights.

 None
bias_distribution Optional[GaussianDistribution]

The distribution for the bias.

 None
Source code in illia/nn/torch/conv1d.py 
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
def __init__(
    self,
    input_channels: int,
    output_channels: int,
    kernel_size: int,
    stride: int = 1,
    padding: int = 0,
    dilation: int = 1,
    groups: int = 1,
    weights_distribution: Optional[GaussianDistribution] = None,
    bias_distribution: Optional[GaussianDistribution] = None,
    **kwargs: Any,
) -> None:
    """
    Definition of a Bayesian Convolution 1D layer.

    Args:
        input_channels: Number of channels in the input image.
        output_channels: Number of channels produced by the
            convolution.
        kernel_size: Size of the convolving kernel.
        stride: Stride of the convolution. Deafults to 1.
        padding: Padding added to all four sides of the input.
            Defaults to 0.
        dilation: Spacing between kernel elements.
        groups: Number of blocked connections from input channels
            to output channels. Defaults to 1.
        weights_distribution: The distribution for the weights.
        bias_distribution: The distribution for the bias.
    """

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

    # Set attributes
    self.conv_params: tuple[int, ...] = (stride, padding, dilation, groups)

    # Set weights distribution
    if weights_distribution is None:
        # Define weights distribution
        self.weights_distribution = GaussianDistribution(
            (output_channels, input_channels // groups, kernel_size)
        )
    else:
        self.weights_distribution = weights_distribution

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

    # Sample initial weights
    weights = self.weights_distribution.sample()
    bias = self.bias_distribution.sample()

    # Register buffers
    self.register_buffer("weights", weights)
    self.register_buffer("bias", bias)

6.1.1 forward(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/torch/conv1d.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
def forward(self, inputs: torch.Tensor) -> torch.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].
    """

    # Forward depending of frozen state
    if not self.frozen:
        self.weights = self.weights_distribution.sample()
        self.bias = self.bias_distribution.sample()
    elif self.weights is None or self.bias is None:
        raise ValueError("Module has been frozen with undefined weights")

    # Execute torch forward
    # pylint: disable=E1102
    return F.conv1d(inputs, self.weights, self.bias, *self.conv_params)

6.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/torch/conv1d.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
@torch.jit.export
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.weights is None:
        self.weights = self.weights_distribution.sample()

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

    # Detach weights and bias
    self.weights = self.weights.detach()
    self.bias = self.bias.detach()

6.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/torch/conv1d.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
@torch.jit.export
def kl_cost(self) -> tuple[torch.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: torch.Tensor = self.weights_distribution.log_prob(
        self.weights
    ) + self.bias_distribution.log_prob(self.bias)

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

    return log_probs, num_params

This module contains the code for the bayesian Conv2d.

6.2 Conv2d(input_channels, output_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, weights_distribution=None, bias_distribution=None, **kwargs)

This class is the bayesian implementation of the Conv2d class.

Definition of a Bayesian Convolution 2D layer.

Parameters:

Name Type Description Default
kernel_size int | tuple[int, int]

Size of the convolving kernel.

 required
stride int | tuple[int, int]

Stride of the convolution. Deafults to 1.

 1
padding int | tuple[int, int]

Padding added to all four sides of the input.

 0
dilation int | tuple[int, int]

Spacing between kernel elements.

 1
groups int

Number of blocked connections from input channels to output channels.

 1
weights_distribution Optional[GaussianDistribution]

The distribution for the weights.

 None
bias_distribution Optional[GaussianDistribution]

The distribution for the bias.

 None
Source code in illia/nn/torch/conv2d.py 
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
def __init__(
    self,
    input_channels: int,
    output_channels: int,
    kernel_size: int | tuple[int, int],
    stride: int | tuple[int, int] = 1,
    padding: int | tuple[int, int] = 0,
    dilation: int | tuple[int, int] = 1,
    groups: int = 1,
    weights_distribution: Optional[GaussianDistribution] = None,
    bias_distribution: Optional[GaussianDistribution] = None,
    **kwargs: Any,
) -> None:
    """
    Definition of a Bayesian Convolution 2D layer.

    Args:
        kernel_size: Size of the convolving kernel.
        stride: Stride of the convolution. Deafults to 1.
        padding: Padding added to all four sides of the input.
        dilation: Spacing between kernel elements.
        groups: Number of blocked connections from input channels
            to output channels.
        weights_distribution: The distribution for the weights.
        bias_distribution: The distribution for the bias.
    """

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

    # 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

    # Set weights distribution
    if weights_distribution is None:
        # Extend kernel if we only have 1 value
        if isinstance(self.kernel_size, int):
            self.kernel_size = (self.kernel_size, self.kernel_size)

        # Define weights distribution
        self.weights_distribution = GaussianDistribution(
            (
                self.output_channels,
                self.input_channels // self.groups,
                *self.kernel_size,
            )
        )
    else:
        self.weights_distribution = weights_distribution

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

    # Sample initial weights
    weights = self.weights_distribution.sample()
    bias = self.bias_distribution.sample()

    # Register buffers
    self.register_buffer("weights", weights)
    self.register_buffer("bias", bias)

6.2.1 forward(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/torch/conv2d.py 
141
142
143
144
145
146
147
148
149
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
def forward(self, inputs: torch.Tensor) -> torch.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].
    """

    # Forward depending of frozen state
    if not self.frozen:
        self.weights = self.weights_distribution.sample()
        self.bias = self.bias_distribution.sample()
    else:
        if self.weights is None or self.bias is None:
            self.weights = self.weights_distribution.sample()
            self.bias = self.bias_distribution.sample()

    # Execute torch forward
    # pylint: disable=E1102
    return F.conv2d(
        input=inputs,
        weight=self.weights,
        bias=self.bias,
        stride=self.stride,
        padding=self.padding,
        dilation=self.dilation,
        groups=self.groups,
    )

6.2.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/torch/conv2d.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
@torch.jit.export
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.weights is None:
        self.weights = self.weights_distribution.sample()

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

    # Detach weights and bias
    self.weights = self.weights.detach()
    self.bias = self.bias.detach()

6.2.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/torch/conv2d.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
@torch.jit.export
def kl_cost(self) -> tuple[torch.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: torch.Tensor = self.weights_distribution.log_prob(
        self.weights
    ) + self.bias_distribution.log_prob(self.bias)

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

    return log_probs, num_params

This module contains the code for bayesian Embedding layer.

6.3 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]

distribution for the weights of the layer.

 None
Source code in illia/nn/torch/embedding.py 
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
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: distribution for the weights of the
            layer.
    """

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

    # Set embeddings atributtes
    self.embedding_params: tuple[Any, ...] = (
        padding_idx,
        max_norm,
        norm_type,
        scale_grad_by_freq,
        sparse,
    )

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

    # Sample initial weights
    weights = self.weights_distribution.sample()

    # Register buffers
    self.register_buffer("weights", weights)

6.3.1 forward(inputs)

This method is the forward pass of the layer.

Parameters:

Name Type Description Default
inputs Tensor

input tensor. Dimensions: [*].

 required

Raises:

Type Description
ValueError

Module has been frozen with undefined weights.

Returns:

Type Description
Tensor

outputs tensor. Dimension: [*, embedding dim].
Source code in illia/nn/torch/embedding.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
    """
    This method is the forward pass of the layer.

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

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

    Returns:
        outputs tensor. Dimension: [*, embedding dim].
    """

    # Forward depeding of frozen state
    if not self.frozen:
        self.weights = self.weights_distribution.sample()
    elif self.weights is None:
        raise ValueError("Module has been frozen with undefined weights")

    # Run torch forward
    outputs: torch.Tensor = F.embedding(
        inputs, self.weights, *self.embedding_params
    )

    return outputs

6.3.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/torch/embedding.py 
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
@torch.jit.export
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.weights is None:
        self.weights = self.weights_distribution.sample()

    # detach weights
    self.weights = self.weights.detach()

6.3.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/torch/embedding.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
@torch.jit.export
def kl_cost(self) -> tuple[torch.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.
    """

    # get log posterior and log prior
    log_probs: torch.Tensor = self.weights_distribution.log_prob(self.weights)

    # get number of parameters
    num_params: int = self.weights_distribution.num_params()

    return log_probs, num_params

This module contains the code for Linear Bayesian layer.

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

This class is the bayesian implementation of the torch Linear layer.

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]

GaussianDistribution for the weights of the layer. Defaults to None.

 None
bias_distribution Optional[GaussianDistribution]

GaussianDistribution for the bias of the layer. Defaults to None.

 None
Source code in illia/nn/torch/linear.py 
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
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: GaussianDistribution for the weights of the
            layer. Defaults to None.
        bias_distribution: GaussianDistribution for the bias of the layer.
            Defaults to None.
    """

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

    # 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

    # Sample initial weights
    weights = self.weights_distribution.sample()
    bias = self.bias_distribution.sample()

    # Register buffers
    self.register_buffer("weights", weights)
    self.register_buffer("bias", bias)

6.4.1 forward(inputs)

This method is the forward pass of the layer.

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

outputs tensor. Dimensions: [batch, *].
Source code in illia/nn/torch/linear.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
    """
    This method is the forward pass of the layer.

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

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

    Returns:
        outputs tensor. Dimensions: [batch, *].
    """

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

    # Compute outputs
    # pylint: disable=E1102
    outputs: torch.Tensor = F.linear(inputs, self.weights, self.bias)

    return outputs

6.4.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/torch/linear.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
@torch.jit.export
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.weights is None:
        self.weights = self.weights_distribution.sample()

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

    # Detach weights and bias
    self.weights = self.weights.detach()
    self.bias = self.bias.detach()

6.4.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/torch/linear.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
@torch.jit.export
def kl_cost(self) -> tuple[torch.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: torch.Tensor = self.weights_distribution.log_prob(
        self.weights
    ) + self.bias_distribution.log_prob(self.bias)

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

    return log_probs, num_params

This module contains the code for the bayesian LSTM.

6.5 LSTM(num_embeddings, embeddings_dim, hidden_size, output_size, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False, **kwargs)

This class is the bayesian implementation of the torch LSTM layer.

Source code in illia/nn/torch/lstm.py 
 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
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
134
135
136
137
138
139
140
141
def __init__(
    self,
    num_embeddings: int,
    embeddings_dim: int,
    hidden_size: int,
    output_size: 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,
    **kwargs: Any,
) -> None:

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

    # Set attributes
    self.num_embeddings = num_embeddings
    self.embeddings_dim = embeddings_dim
    self.hidden_size = hidden_size
    self.output_size = output_size
    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

    # Define the Embedding layer
    self.embedding = Embedding(
        num_embeddings=self.num_embeddings,
        embeddings_dim=self.embeddings_dim,
        padding_idx=self.padding_idx,
        max_norm=self.max_norm,
        norm_type=self.norm_type,
        scale_grad_by_freq=self.scale_grad_by_freq,
        sparse=self.sparse,
    )

    # Initialize weights
    # Forget gate
    self.wf_distribution = GaussianDistribution(
        (self.hidden_size, self.embeddings_dim + self.hidden_size)
    )
    self.bf_distribution = GaussianDistribution((self.hidden_size,))

    # Input gate
    self.wi_distribution = GaussianDistribution(
        (self.hidden_size, self.embeddings_dim + self.hidden_size)
    )
    self.bi_distribution = GaussianDistribution((self.hidden_size,))

    # Candidate gate
    self.wc_distribution = GaussianDistribution(
        (self.hidden_size, self.embeddings_dim + self.hidden_size)
    )
    self.bc_distribution = GaussianDistribution((self.hidden_size,))

    # Output gate
    self.wo_distribution = GaussianDistribution(
        (self.hidden_size, self.embeddings_dim + self.hidden_size)
    )
    self.bo_distribution = GaussianDistribution((self.hidden_size,))

    # Final gate
    self.wv_distribution = GaussianDistribution(
        (self.output_size, self.hidden_size)
    )
    self.bv_distribution = GaussianDistribution((self.output_size,))

    # Sample initial weights and register buffers
    # Forget gate
    wf = self.wf_distribution.sample()
    bf = self.bf_distribution.sample()
    self.register_buffer("wf", wf)
    self.register_buffer("bf", bf)

    # Input gate
    wi = self.wi_distribution.sample()
    bi = self.bi_distribution.sample()
    self.register_buffer("wi", wi)
    self.register_buffer("bi", bi)

    # Candidate gate
    wc = self.wc_distribution.sample()
    bc = self.bc_distribution.sample()
    self.register_buffer("wc", wc)
    self.register_buffer("bc", bc)

    # Output gate
    wo = self.wo_distribution.sample()
    bo = self.bo_distribution.sample()
    self.register_buffer("wo", wo)
    self.register_buffer("bo", bo)

    # Final output layer
    wv = self.wv_distribution.sample()
    bv = self.bv_distribution.sample()
    self.register_buffer("wv", wv)
    self.register_buffer("bv", bv)

6.5.1 forward(inputs, init_states=None)

Performs a forward pass through the Bayesian LSTM 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
tuple[Tensor, tuple[Tensor, Tensor]]

Output tensor after passing through the layer. Dimensions: [batch, output channels, output width, output height].
Source code in illia/nn/torch/lstm.py 
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
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
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
def forward(
    self,
    inputs: torch.Tensor,
    init_states: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    """
    Performs a forward pass through the Bayesian LSTM 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].
    """

    # Sample weights if not frozen
    if not self.frozen:
        self.wf = self.wf_distribution.sample()
        self.bf = self.bf_distribution.sample()
        self.wi = self.wi_distribution.sample()
        self.bi = self.bi_distribution.sample()
        self.wc = self.wc_distribution.sample()
        self.bc = self.bc_distribution.sample()
        self.wo = self.wo_distribution.sample()
        self.bo = self.bo_distribution.sample()
        self.wv = self.wv_distribution.sample()
        self.bv = self.bv_distribution.sample()
    else:
        if any(w is None for w in [self.wf, self.wi, self.wc, self.wo, self.wv]):
            self.wf = self.wf_distribution.sample()
            self.bf = self.bf_distribution.sample()
            self.wi = self.wi_distribution.sample()
            self.bi = self.bi_distribution.sample()
            self.wc = self.wc_distribution.sample()
            self.bc = self.bc_distribution.sample()
            self.wo = self.wo_distribution.sample()
            self.bo = self.bo_distribution.sample()
            self.wv = self.wv_distribution.sample()
            self.bv = self.bv_distribution.sample()

    # Apply embedding layer to input indices
    inputs = inputs.squeeze(dim=-1)
    inputs = self.embedding(inputs)
    batch_size, seq_len, _ = inputs.size()

    # Initialize h_t and c_t if init_states is None
    if init_states is None:
        device = inputs.device
        h_t = torch.zeros(batch_size, self.hidden_size, device=device)
        c_t = torch.zeros(batch_size, self.hidden_size, device=device)
    else:
        h_t, c_t = init_states[0], init_states[1]

    for t in range(seq_len):
        # Shape: (batch_size, embedding_dim)
        x_t = inputs[:, t, :]

        # Concatenate input and hidden state
        # Shape: (batch_size, embedding_dim + hidden_size)
        z_t = torch.cat([x_t, h_t], dim=1)

        # Forget gate
        ft = torch.sigmoid(z_t @ self.wf.t() + self.bf)

        # Input gate
        it = torch.sigmoid(z_t @ self.wi.t() + self.bi)

        # Candidate cell state
        can = torch.tanh(z_t @ self.wc.t() + self.bc)

        # Output gate
        ot = torch.sigmoid(z_t @ self.wo.t() + self.bo)

        # Update cell state
        c_t = c_t * ft + can * it

        # Update hidden state
        h_t = ot * torch.tanh(c_t)

    # Compute final output
    y_t = h_t @ self.wv.t() + self.bv

    return y_t, (h_t, c_t)

6.5.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/torch/lstm.py 
143
144
145
146
147
148
149
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
194
@torch.jit.export
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

    # Freeze embedding layer
    self.embedding.freeze()

    # Forget gate
    if self.wf is None:
        self.wf = self.wf_distribution.sample()
    if self.bf is None:
        self.bf = self.bf_distribution.sample()
    self.wf = self.wf.detach()
    self.bf = self.bf.detach()

    # Input gate
    if self.wi is None:
        self.wi = self.wi_distribution.sample()
    if self.bi is None:
        self.bi = self.bi_distribution.sample()
    self.wi = self.wi.detach()
    self.bi = self.bi.detach()

    # Candidate gate
    if self.wc is None:
        self.wc = self.wc_distribution.sample()
    if self.bc is None:
        self.bc = self.bc_distribution.sample()
    self.wc = self.wc.detach()
    self.bc = self.bc.detach()

    # Output gate
    if self.wo is None:
        self.wo = self.wo_distribution.sample()
    if self.bo is None:
        self.bo = self.bo_distribution.sample()
    self.wo = self.wo.detach()
    self.bo = self.bo.detach()

    # Final output layer
    if self.wv is None:
        self.wv = self.wv_distribution.sample()
    if self.bv is None:
        self.bv = self.bv_distribution.sample()
    self.wv = self.wv.detach()
    self.bv = self.bv.detach()

6.5.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/torch/lstm.py 
196
197
198
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
234
235
236
237
238
239
240
241
242
243
244
245
246
@torch.jit.export
def kl_cost(self) -> tuple[torch.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 for each pair of weights and bias
    # Forget gate
    log_probs_f: torch.Tensor = self.wf_distribution.log_prob(
        self.wf
    ) + self.bf_distribution.log_prob(self.bf)
    # Input gate
    log_probs_i: torch.Tensor = self.wi_distribution.log_prob(
        self.wi
    ) + self.bi_distribution.log_prob(self.bi)
    # Candidate gate
    log_probs_c: torch.Tensor = self.wc_distribution.log_prob(
        self.wc
    ) + self.bc_distribution.log_prob(self.bc)
    # Output gate
    log_probs_o: torch.Tensor = self.wo_distribution.log_prob(
        self.wo
    ) + self.bo_distribution.log_prob(self.bo)
    # Final output layer
    log_probs_v: torch.Tensor = self.wv_distribution.log_prob(
        self.wv
    ) + self.bv_distribution.log_prob(self.bv)

    # Compute the total loss
    log_probs = log_probs_f + log_probs_i + log_probs_c + log_probs_o + log_probs_v

    # Compute number of parameters
    num_params: int = (
        self.wf_distribution.num_params()
        + self.bf_distribution.num_params()
        + self.wi_distribution.num_params()
        + self.bi_distribution.num_params()
        + self.wc_distribution.num_params()
        + self.bc_distribution.num_params()
        + self.wo_distribution.num_params()
        + self.bo_distribution.num_params()
        + self.wv_distribution.num_params()
        + self.bv_distribution.num_params()
    )

    return log_probs, num_params