PyDLLib 1.0__py3-none-any.whl

DLL/DeepLearning/Layers/Regularisation/_Dropout.py

@@ -0,0 +1,79 @@
+import torch
+
+from ._BaseRegularisation import BaseRegularisation
+
+
+class Dropout(BaseRegularisation):
+    """
+    The dropout layer for neural networks.
+
+    Args:
+        p (float, optional): The probability of a node being dropped out. Must be strictly between 0 and 1. Defaults to 0.5.
+    """
+    def __init__(self, p=0.5, **kwargs):
+        if p <= 0 or 1 <= p:
+            raise ValueError("p must be strictly between 0 and 1.")
+        
+        super().__init__(**kwargs)
+        self.p = 1 - p
+        self.name = "Dropout"
+    
+    def initialise_layer(self, **kwargs):
+        """
+        :meta private:
+        """
+
+        super().initialise_layer(**kwargs)
+
+    def forward(self, input, training=False, **kwargs):
+        """
+        Sets some values of the input to zero with probability p.
+
+        Args:
+            input (torch.Tensor of shape (batch_size, channels, ...)): The input to the layer. Must be a torch.Tensor of the spesified shape given by layer.input_shape.
+            training (bool, optional): The boolean flag deciding if the model is in training mode. Defaults to False.
+
+        Returns:
+            torch.Tensor: The output tensor after the transformation with the same shape as the input.
+        """
+        if not isinstance(input, torch.Tensor):
+            raise TypeError("input must be a torch.Tensor.")
+        if input.shape[-len(self.input_shape):] != self.input_shape:
+            raise ValueError(f"input is not the same shape as the spesified input_shape ({input.shape[-len(self.input_shape):], self.input_shape}).")
+        if not isinstance(training, bool):
+            raise TypeError("training must be a boolean.")
+        
+        self.input = input
+        if training:
+            self.mask = torch.rand(size=self.input.shape, dtype=self.data_type, device=self.device) < self.p
+            self.output = self.input * self.mask / self.p
+        else:
+            self.output = self.input
+        return self.output
+    
+    def backward(self, dCdy, **kwargs):
+        """
+        Calculates the gradient of the loss function with respect to the input of the layer.
+
+        Args:
+            dCdy (torch.Tensor of the same shape as returned from the forward method): The gradient given by the next layer.
+            
+        Returns:
+            torch.Tensor of shape (batch_size, channels, ...): The new gradient after backpropagation through the layer.
+        """
+        if not hasattr(self, "mask"):
+            raise ValueError("Forward method should be called before backward.")
+        if not isinstance(dCdy, torch.Tensor):
+            raise TypeError("dCdy must be a torch.Tensor.")
+        if dCdy.shape != self.output.shape:
+            raise ValueError(f"dCdy is not the same shape as the spesified output_shape ({dCdy.shape, self.output.shape}).")
+        
+        dCdx = dCdy * self.mask / self.p
+        return dCdx
+    
+    def summary(self, offset=""):
+        """
+        :meta private:
+        """
+        output_shape = self.output_shape[0] if len(self.output_shape) == 1 else self.output_shape
+        return offset + f"{self.name} - Output: ({output_shape}) - Keep probability: {self.p}"

DLL/DeepLearning/Layers/Regularisation/_GroupNormalisation.py

@@ -0,0 +1,112 @@
+import torch
+
+from ._BaseRegularisation import BaseRegularisation
+
+
+class GroupNorm(BaseRegularisation):
+    """
+    The group normalisation layer for neural networks. Computes the group norm of a batch along axis=1
+
+    Args:
+        num_groups (int, optional): The number of groups used in the normalisation. Must be a positive integer. Defaults to 32. The number of channels must be evenly divisible by num_groups. If is set to 1, is identical to layer normalisation and if batch_size, is identical to the instance normalisation.
+    """
+    def __init__(self, num_groups=32, **kwargs):
+        if num_groups is not None and (not isinstance(num_groups, int) or num_groups <= 0):
+            raise ValueError("num_groups must be a positive integer.")
+
+        super().__init__(**kwargs)
+        self.num_groups = num_groups
+        self.epsilon = 1e-6
+        self.name = "Group normalisation"
+    
+    def initialise_layer(self, **kwargs):
+        """
+        :meta private:
+        """
+        super().initialise_layer(**kwargs)
+
+        if self.output_shape[0] % self.num_groups != 0:
+            raise ValueError("The number of channels must be evenly divisible by num_groups.")
+        
+        self.gamma = torch.ones(self.output_shape, device=self.device, dtype=self.data_type)
+        self.beta = torch.zeros(self.output_shape, device=self.device, dtype=self.data_type)
+        self.nparams = 2 * self.output_shape[0]
+
+    def forward(self, input, **kwargs):
+        """
+        Normalises the input to have zero mean and one variance accross self.num_groups groups accross the channel dimension with the following equation:
+        
+        .. math::
+            y = \\gamma\\frac{x - \\mathbb{E}[x]}{\\sqrt{\\text{var}(x) + \\epsilon}} + \\beta,
+        
+        where :math:`x` is the input, :math:`\\mathbb{E}[x]` is the expected value or the mean accross each group, :math:`\\text{var}(x)` is the variance accross the variance accross each group, :math:`\\epsilon` is a small constant and :math:`\\gamma` and :math:`\\beta` are trainable parameters.
+
+        Args:
+            input (torch.Tensor of shape (batch_size, channels, ...)): The input to the layer. Must be a torch.Tensor of the spesified shape given by layer.input_shape.
+            training (bool, optional): The boolean flag deciding if the model is in training mode. Defaults to False.
+            
+        Returns:
+            torch.Tensor: The output tensor after the normalisation with the same shape as the input.
+        """
+        if not isinstance(input, torch.Tensor):
+            raise TypeError("input must be a torch.Tensor.")
+        if input.shape[1:] != self.input_shape:
+            raise ValueError(f"input is not the same shape as the spesified input_shape ({input.shape[1:], self.input_shape}).")
+
+        elements_per_group = input.shape[1] // self.num_groups
+        self.input = input
+        batch_size = input.shape[0]
+        self.input = input
+        self.input_reshaped = self.input.view(batch_size, self.num_groups, elements_per_group, *input.shape[2:])
+        mean = 1.0 / elements_per_group * self.input_reshaped.sum(2, keepdim=True)
+
+        self.x_centered = self.input_reshaped - mean
+        self.x_centered_squared = self.x_centered ** 2
+        self.var = 1.0 / elements_per_group * self.x_centered_squared.sum(2, keepdim=True)  # biased variance
+
+        self.inv_std = (self.var + self.epsilon) ** -0.5
+        self.x_norm = self.x_centered * self.inv_std
+        self.x_reshaped = self.x_norm.view(self.input.shape)
+        self.output = self.x_reshaped * self.gamma + self.beta
+        return self.output
+    
+    def backward(self, dCdy, **kwargs):
+        """
+        Calculates the gradient of the loss function with respect to the input of the layer. Also calculates the gradients of the loss function with respect to the model parameters.
+
+        Args:
+            dCdy (torch.Tensor of the same shape as returned from the forward method): The gradient given by the next layer.
+            
+        Returns:
+            torch.Tensor of shape (batch_size, channels, ...): The new gradient after backpropagation through the layer.
+        """
+        if not isinstance(dCdy, torch.Tensor):
+            raise TypeError("dCdy must be a torch.Tensor.")
+        if dCdy.shape[1:] != self.output.shape[1:]:
+            raise ValueError(f"dCdy is not the same shape as the spesified output_shape ({dCdy.shape[1:], self.output.shape[1:]}).")
+        
+        batch_size = dCdy.shape[0]
+        elements_per_group = self.output.shape[1] // self.num_groups
+        dCdx_reshaped = dCdy * self.gamma
+        dCdgamma = (dCdy * self.x_reshaped).mean(axis=0)
+        dCdbeta = dCdy.mean(axis=0)
+        self.gamma.grad += dCdgamma
+        self.beta.grad += dCdbeta
+
+        dCdx_norm = dCdx_reshaped.view(batch_size, self.num_groups, elements_per_group, *self.output.shape[2:])
+        dCdx_centered = dCdx_norm * self.inv_std
+        dCdinv_std = (dCdx_norm * self.x_centered).sum(2, keepdim=True)
+        dCdvar = -0.5 * ((self.var + self.epsilon) ** -1.5) * dCdinv_std
+        dCdx_centered_squared = 1.0 / elements_per_group * torch.ones_like(self.x_centered_squared, device=self.device, dtype=self.data_type) * dCdvar
+        dCdx_centered += 2 * self.x_centered * dCdx_centered_squared
+        dCdinput_reshaped = dCdx_centered.clone()
+        dCdmean = -(dCdx_centered).sum(2, keepdim=True)
+        dCdinput_reshaped += 1.0 / elements_per_group * torch.ones_like(self.input_reshaped, device=self.device, dtype=self.data_type) * dCdmean
+        dCdx = dCdinput_reshaped.view(self.output.shape)
+        return dCdx.view(self.output.shape)
+
+    def get_parameters(self):
+        """
+        :meta private:
+        """
+        return (self.gamma, self.beta)

DLL/DeepLearning/Layers/Regularisation/_InstanceNormalisation.py

@@ -0,0 +1,50 @@
+from ._GroupNormalisation import GroupNorm
+
+
+class InstanceNorm(GroupNorm):
+    """
+    The instance normalisation layer for neural networks. Computes the group norm of a batch along axis=1 with the same number of groups as channels.
+    """
+    def __init__(self, **kwargs):
+        super().__init__(None, **kwargs)
+        self.name = "Instance normalisation"
+
+    def initialise_layer(self, **kwargs):
+        """
+        :meta private:
+        """
+        if "input_shape" not in kwargs or not isinstance(kwargs.get("input_shape"), tuple | list):
+            raise ValueError(f"input_shape must be passed as a key word argument to layer.initialise_layer as a tuple or a list. Currently {kwargs.get('input_shape', None)}.")
+
+        self.num_groups = kwargs.get("input_shape")[0]
+        super().initialise_layer(**kwargs)
+
+    def forward(self, input, **kwargs):
+        """
+        Normalises the input to have zero mean and one variance accross the channel dimension with the following equation:
+        
+        .. math::
+            y = \\gamma\\frac{x - \\mathbb{E}[x]}{\\sqrt{\\text{var}(x) + \\epsilon}} + \\beta,
+        
+        where :math:`x` is the input, :math:`\\mathbb{E}[x]` is the expected value or the mean accross the channel dimension, :math:`\\text{var}(x)` is the variance accross the variance accross the channel dimension, :math:`\\epsilon` is a small constant and :math:`\\gamma` and :math:`\\beta` are trainable parameters.
+
+        Args:
+            input (torch.Tensor of shape (batch_size, channels, ...)): The input to the layer. Must be a torch.Tensor of the spesified shape given by layer.input_shape.
+
+        Returns:
+            torch.Tensor: The output tensor after the normalisation with the same shape as the input.
+        """
+
+        return super().forward(input, **kwargs)
+    
+    def backward(self, dCdy, **kwargs):
+        """
+        Calculates the gradient of the loss function with respect to the input of the layer. Also calculates the gradients of the loss function with respect to the model parameters.
+
+        Args:
+            dCdy (torch.Tensor of the same shape as returned from the forward method): The gradient given by the next layer.
+            
+        Returns:
+            torch.Tensor of shape (batch_size, channels, ...): The new gradient after backpropagation through the layer.
+        """
+        return super().backward(dCdy, **kwargs)

DLL/DeepLearning/Layers/Regularisation/_LayerNormalisation.py

@@ -0,0 +1,37 @@
+from ._GroupNormalisation import GroupNorm
+
+
+class LayerNorm(GroupNorm):
+    def __init__(self, **kwargs):
+        super().__init__(1, **kwargs)
+        self.name = "Layer normalisation"
+
+    def forward(self, input, **kwargs):
+        """
+        Normalises the input to have zero mean and one variance accross the channel dimension with the following equation:
+        
+        .. math::
+            y = \\gamma\\frac{x - \\mathbb{E}[x]}{\\sqrt{\\text{var}(x) + \\epsilon}} + \\beta,
+        
+        where :math:`x` is the input, :math:`\\mathbb{E}[x]` is the expected value or the mean accross the channel dimension, :math:`\\text{var}(x)` is the variance accross the variance accross the channel dimension, :math:`\\epsilon` is a small constant and :math:`\\gamma` and :math:`\\beta` are trainable parameters.
+
+        Args:
+            input (torch.Tensor of shape (batch_size, channels, ...)): The input to the layer. Must be a torch.Tensor of the spesified shape given by layer.input_shape.
+
+        Returns:
+            torch.Tensor: The output tensor after the normalisation with the same shape as the input.
+        """
+
+        return super().forward(input, **kwargs)
+    
+    def backward(self, dCdy, **kwargs):
+        """
+        Calculates the gradient of the loss function with respect to the input of the layer. Also calculates the gradients of the loss function with respect to the model parameters.
+
+        Args:
+            dCdy (torch.Tensor of the same shape as returned from the forward method): The gradient given by the next layer.
+            
+        Returns:
+            torch.Tensor of shape (batch_size, channels, ...): The new gradient after backpropagation through the layer.
+        """
+        return super().backward(dCdy, **kwargs)

DLL/DeepLearning/Layers/Regularisation/__init__.py

@@ -0,0 +1,5 @@
+from ._BatchNormalisation import BatchNorm
+from ._Dropout import Dropout
+from ._GroupNormalisation import GroupNorm
+from ._InstanceNormalisation import InstanceNorm
+from ._LayerNormalisation import LayerNorm

DLL/DeepLearning/Layers/_Add.py

@@ -0,0 +1,109 @@
+import torch
+
+from ._BaseLayer import BaseLayer
+from ...Exceptions import NotCompiledError
+from .Activations._Activation import Activation
+from .Regularisation._BaseRegularisation import BaseRegularisation
+
+
+class Add(BaseLayer):
+    """
+    The addition layer.
+
+    Args:
+        layer1 (DLL.DeepLearning.Layers.BaseLayer object): The first layer the input is passed to. The results of each layer are added together. The input and outpput shapes of the layers must be the same.
+        layer2 (DLL.DeepLearning.Layers.BaseLayer object): The second layer the input is passed to. The results of each layer are added together. The input and outpput shapes of the layers must be the same.
+    """
+    def __init__(self, layer1, layer2, activation=None, normalisation=None, **kwargs):
+        if not isinstance(layer1, BaseLayer) or not isinstance(layer2, BaseLayer):
+            raise TypeError("layers must be an instances of DLL.DeepLearning.Layers.BaseLayer")
+        if layer1.output_shape != layer2.output_shape and layer1.output_shape is not None and layer2.output_shape is not None:
+            raise ValueError("Layers must have the same output shape.")
+        if not isinstance(activation, Activation) and activation is not None:
+            raise ValueError("activation must be from DLL.DeepLearning.Layers.Activations or None.")
+        if not isinstance(normalisation, BaseRegularisation) and normalisation is not None:
+            raise ValueError("normalisation must be from DLL.DeepLearning.Layers.Regularisation or None.")
+
+        super().__init__(layer1.output_shape, activation=activation, normalisation=normalisation, **kwargs)
+        self.name = "Add"
+        self.layer1 = layer1
+        self.layer2 = layer2
+
+    def initialise_layer(self, input_shape, data_type, device):
+        """
+        :meta private:
+        """
+        if self.layer1.output_shape is None and input_shape != self.layer2.output_shape:
+            raise ValueError(f"Layers must have the same output shape {input_shape} vs {self.layer2.output_shape}.")
+        if self.layer2.output_shape is None and input_shape != self.layer1.output_shape:
+            raise ValueError(f"Layers must have the same output shape {input_shape} vs {self.layer1.output_shape}.")
+
+        self.layer1.initialise_layer(input_shape, data_type, device)
+        self.layer2.initialise_layer(input_shape, data_type, device)
+        super().initialise_layer(input_shape, data_type, device)
+
+    def forward(self, input, training=False, **kwargs):
+        """
+        Computes the forward values of the input layers and adds them together.
+
+        Args:
+            input (torch.Tensor): The input to the layer. Must be a torch.Tensor of the spesified shape given by layer.input_shape.
+            training (bool, optional): The boolean flag deciding if the model is in training mode. Defaults to False.
+            
+        Returns:
+            torch.Tensor: The output tensor after the transformations with the spesified shape.
+        """
+        if not isinstance(input, torch.Tensor):
+            raise TypeError("input must be a torch.Tensor.")
+        if input.shape[1:] != self.input_shape:
+            raise ValueError(f"Input shape {input.shape[1:]} does not match the expected shape {self.input_shape}.")
+        if not isinstance(training, bool):
+            raise TypeError("training must be a boolean.")
+
+        val1 = self.layer1.forward(input, training=training)
+        val2 = self.layer2.forward(input.flip(1), training=training)
+        self.output = val1 + val2
+        if self.normalisation: self.output = self.normalisation.forward(self.output, training=training)
+        if self.activation: self.output = self.activation.forward(self.output)
+        return self.output
+
+    def backward(self, dCdy, **kwargs):
+        """
+        Calculates the gradient of the loss function with respect to the input of the layers. Also calculates the gradients of the loss function with respect to the model parameters.
+
+        Args:
+            dCdy (torch.Tensor of the same shape as returned from the forward method): The gradient given by the next layer.
+
+        Returns:
+            torch.Tensor of the spesified shape: The new gradient after backpropagation through the layer.
+        """
+        if not isinstance(dCdy, torch.Tensor):
+            raise TypeError("dCdy must be a torch.Tensor.")
+        if dCdy.shape[1:] != self.output.shape[1:]:
+            raise ValueError(f"dCdy is not the same shape as the spesified output_shape ({dCdy.shape[1:], self.output.shape[1:]}).")
+        
+        if self.activation: dCdy = self.activation.backward(dCdy)
+        if self.normalisation: dCdy = self.normalisation.backward(dCdy)
+        dCdx1 = self.layer1.backward(dCdy, **kwargs)
+        dCdx2 = self.layer2.backward(dCdy, **kwargs)
+
+        dCdx = dCdx1 + dCdx2
+        return dCdx
+
+    def get_parameters(self):
+        """
+        :meta private:
+        """
+        return (*self.layer1.get_parameters(), *self.layer2.get_parameters(), *super().get_parameters())
+    
+    def get_nparams(self):
+        return self.layer1.get_nparams() + self.layer2.get_nparams()
+    
+    def summary(self, offset=""):
+        if not hasattr(self, "input_shape"):
+            raise NotCompiledError("layer must be initialized correctly before calling layer.summary().")
+
+        sublayer_offset = offset + "    "
+        summary1 = "\n" + self.layer1.summary(sublayer_offset)
+        summary2 = "\n" + self.layer2.summary(sublayer_offset)
+        return super().summary(offset) + summary1 + summary2

DLL/DeepLearning/Layers/_Attention.py

@@ -0,0 +1,161 @@
+import torch
+
+from ._BaseLayer import BaseLayer
+from . import Dense
+from .Activations import SoftMax
+from .Regularisation import Dropout
+from .Activations._Activation import Activation
+from .Regularisation._BaseRegularisation import BaseRegularisation
+
+
+class MultiHeadAttention(BaseLayer):
+    """
+    The multi head attention layer.
+
+    Args:
+        output_shape (tuple[int]): The output_shape of the model not containing the batch_size. Must be a tuple of positive integers. The returned tensor is of shape (n_samples, seq_len, output_shape) if len(output_shape) == 2 else (n_samples, seq_len).
+        n_heads (int): The number of heads used in the layer. The output dimension must be divisible by n_heads.
+        use_mask (bool): Determines if a mask is used to make the model only consider past tokens. Must be a boolean.
+        dropout (float): The probability of a node being dropped out. Must be in range [0, 1). Defaults to 0.0.
+        activation (:ref:`activations_section_label` | None, optional): The activation used after this layer. If is set to None, no activation is used. Defaults to None. If both activation and regularisation is used, the regularisation is performed first in the forward propagation.
+        normalisation (:ref:`regularisation_layers_section_label` | None, optional): The regularisation layer used after this layer. If is set to None, no regularisation is used. Defaults to None. If both activation and regularisation is used, the regularisation is performed first in the forward propagation.
+    """
+    def __init__(self, output_shape, n_heads=1, use_mask=True, dropout=0.0, activation=None, normalisation=None, **kwargs):
+        if not isinstance(activation, Activation) and activation is not None:
+            raise ValueError("activation must be from DLL.DeepLearning.Layers.Activations or None.")
+        if not isinstance(normalisation, BaseRegularisation) and normalisation is not None:
+            raise ValueError("normalisation must be from DLL.DeepLearning.Layers.Regularisation or None.")
+        if not isinstance(use_mask, bool):
+            raise TypeError("use_mask must be a boolean.")
+        if dropout < 0 or 1 <= dropout:
+            raise ValueError("dropout must be in range [0, 1).")
+        if len(output_shape) == 1 and n_heads != 1 or output_shape[1] % n_heads != 0:
+            raise ValueError("output_dimension must be divisible by n_heads")
+
+        super().__init__(output_shape, activation=activation, normalisation=normalisation, **kwargs)
+        self.name = "MultiHeadAttention"
+        self.use_mask = use_mask
+        self.mask = None
+        self.dropout = dropout
+        self.n_heads = n_heads
+
+    def initialise_layer(self, input_shape, data_type, device):
+        """
+        :meta private:
+        """
+        output_dim = self.output_shape[1] if len(self.output_shape) == 2 else 1
+        if input_shape[1] != output_dim:
+            raise ValueError("input_shape must be the same as the output_shape.")
+        if input_shape[0] != self.output_shape[0]:
+            raise ValueError("Input_shape must have the same seq_len as the output_shape.")
+        
+        super().initialise_layer(input_shape, data_type, device)
+        self.softmax = SoftMax()
+        seq_len = input_shape[0]
+        
+        self.input_linear = Dense((seq_len, 3 * output_dim), bias=False)
+        self.input_linear.initialise_layer(input_shape, data_type, device)
+
+        if self.use_mask:
+            self.mask = torch.tril(torch.ones(seq_len, seq_len)).view(1, 1, seq_len, seq_len) == 0
+        if self.dropout > 0:
+            self.dropout = Dropout(self.dropout)
+            self.dropout.initialise_layer(input_shape=(self.n_heads, seq_len, seq_len), data_type=self.data_type, device=self.device)
+
+        
+    def forward(self, input, training=False, **kwargs):
+        """
+        Applies the attention mechanism on multiple heads.
+
+        .. math::
+        
+            \\begin{align*}
+                y_{\\text{MultiHead}} &= \\text{Concat}(head_1, \\dots, head_{\\text{n_heads}}),\\\\
+                y_{reg} &= f(y_{\\text{MultiHead}}),\\\\
+                y_{activ} &= g(y_{reg}),
+            \\end{align*}
+        
+        where :math:`head_i = \\text{Attention}(Q, K, V) = \\text{softmax}(\\frac{QK^T}{\\sqrt{\\text{output_dim}}})`, :math:`f` is the possible regularisation function and :math:`g` is the possible activation function. :math:`Q, K` and :math:`V` are the query, key and value matricies, which taken from the input by transforming it by a linear layer and splitting the result on the feature axis.
+
+        Args:
+            input (torch.Tensor of shape (n_samples, seq_len, output_shape)): The input to the dense layer. Must be a torch.Tensor of the spesified shape given by layer.input_shape.
+            training (bool, optional): The boolean flag deciding if the model is in training mode. Defaults to False.
+
+        Returns:
+            torch.Tensor of shape (n_samples, seq_len) if output_shape == 0 else (n_samples, seq_len, output_shape): The output tensor after the transformations with the spesified shape.
+        """
+        if not isinstance(input, torch.Tensor):
+            raise TypeError("input must be a torch.Tensor.")
+        if input.shape[1:] != self.input_shape:
+            raise ValueError(f"input is not the same shape as the spesified input_shape ({input.shape[1:], self.input_shape}).")
+        if not isinstance(training, bool):
+            raise TypeError("training must be a boolean.")
+        
+        self.input = input
+        batch_size, seq_len, _ = input.size()
+
+        # split the input to the key, query and the value matricies
+        output_dim = self.output_shape[1] if len(self.output_shape) == 2 else 1
+        input = self.input_linear.forward(input, training=training)  # project the input into key, query and value matricies
+        Q, K, V  = input.chunk(3, dim=2)
+        self.Q = Q.reshape(batch_size, seq_len, self.n_heads, output_dim // self.n_heads).transpose(1, 2)
+        self.K = K.reshape(batch_size, seq_len, self.n_heads, output_dim // self.n_heads).transpose(1, 2)
+        self.V = V.reshape(batch_size, seq_len, self.n_heads, output_dim // self.n_heads).transpose(1, 2)
+
+        output = (self.Q @ self.K.transpose(-2, -1)) * output_dim ** (-0.5)
+        if self.mask is not None: output = output.masked_fill(self.mask, float('-inf'))
+        output = self.softmax.forward(output, training=training)
+        if not isinstance(self.dropout, float | int): output = self.dropout.forward(output, training=training)
+        self.attention_scores = output
+        output = output @ self.V
+        output = output.transpose(1, 2).contiguous().reshape(batch_size, seq_len, output_dim)
+
+        if self.normalisation: output = self.normalisation.forward(output, training=training)
+        if self.activation: output = self.activation.forward(output, training=training)
+        if len(self.output_shape) == 1: output = output.squeeze(dim=2)  # If the output_shape is a 2d tensor, remove the last dimension
+        return output
+
+    def backward(self, dCdy, **kwargs):
+        """
+        Calculates the gradient of the loss function with respect to the input of the layer.
+
+        Args:
+            dCdy (torch.Tensor of shape (n_samples, seq_len, output_dim) if len(output_shape[0]) != 0 else (n_samples, seq_len)): The gradient given by the next layer.
+
+        Returns:
+            torch.Tensor of shape (n_samples, seq_len, output_dim): The new gradient after backpropagation through the layer.
+        """
+        if not isinstance(dCdy, torch.Tensor):
+            raise TypeError("dCdy must be a torch.Tensor.")
+        if dCdy.shape[1:] != self.output_shape:
+            raise ValueError(f"dCdy is not the same shape as the spesified output_shape ({dCdy.shape[1:], self.output_shape}).")
+
+        if len(self.output_shape) == 1: dCdy = dCdy.unsqueeze(dim=2)  # If the output shape was a 2d tensor, add an extra dimension to the end
+        if self.activation: dCdy = self.activation.backward(dCdy)
+        if self.normalisation: dCdy = self.normalisation.backward(dCdy)
+
+        batch_size, seq_len, output_dim = dCdy.size()
+        dCdy = dCdy.reshape(batch_size, seq_len, self.n_heads, output_dim // self.n_heads).transpose(1, 2)
+        dCdV = self.attention_scores.transpose(-1, -2) @ dCdy
+        dCdy = dCdy @ self.V.transpose(-1, -2)
+        if not isinstance(self.dropout, float | int): dCdy = self.dropout.backward(dCdy)
+        dCdy = self.softmax.backward(dCdy)
+        if self.mask is not None: dCdy = dCdy * (~self.mask)
+
+        dCdy = dCdy * output_dim ** (-0.5)
+        dCdQ = dCdy @ self.K
+        dCdK = dCdy.transpose(-2, -1) @ self.Q
+
+        dCdQ = dCdQ.transpose(1, 2).reshape(batch_size, seq_len, output_dim)
+        dCdK = dCdK.transpose(1, 2).reshape(batch_size, seq_len, output_dim)
+        dCdV = dCdV.transpose(1, 2).reshape(batch_size, seq_len, output_dim)
+        dCdx = torch.cat([dCdQ, dCdK, dCdV], dim=2)
+        dCdx = self.input_linear.backward(dCdx)
+
+        return dCdx
+    
+    def get_parameters(self):
+        """
+        :meta private:
+        """
+        return (*self.input_linear.get_parameters(), *super().get_parameters())

DLL/DeepLearning/Layers/_BaseLayer.py

@@ -0,0 +1,87 @@
+import torch
+from copy import deepcopy
+
+from ...Exceptions import NotCompiledError
+
+
+class BaseLayer:
+    def __init__(self, output_shape, input_shape=None, activation=None, normalisation=None, data_type=torch.float32, device=torch.device("cpu")):
+        if output_shape == 0: output_shape = tuple()
+        output_shape = (output_shape,) if isinstance(output_shape, int) else output_shape
+        if not isinstance(output_shape, tuple) and output_shape is not None:
+            raise TypeError(f"output_shape must be a tuple. Currently {output_shape}.")
+        if not isinstance(input_shape, tuple) and input_shape is not None:
+            raise ValueError(f"input_shape must be a tuple or None. Currently {input_shape}.")
+        if not isinstance(data_type, torch.dtype):
+            raise TypeError(f"data_type must be an instance of torch.dtype.")
+        if not isinstance(device, torch.device):
+            raise ValueError('device must be one of torch.device("cpu") or torch.device("cuda").')
+
+        self.input_shape = input_shape
+        self.output_shape = output_shape
+        self.output = None
+        self.input = None
+        self.nparams = 0
+        self.name = "base"
+        self.activation = activation
+        self.normalisation = normalisation
+        self.device = device
+        self.data_type = data_type
+
+    def initialise_layer(self, input_shape, data_type, device):
+        if not isinstance(input_shape, tuple):
+            raise ValueError(f"input_shape must be a tuple or None. Currently {input_shape}.")
+        if not isinstance(data_type, torch.dtype):
+            raise TypeError("data_type must be an instance of torch.dtype.")
+        if not isinstance(data_type, torch.dtype):
+            raise TypeError("data_type must be an instance of torch.dtype.")
+
+        self.input_shape = input_shape
+        self.data_type = data_type
+        self.device = device
+
+        if self.activation:
+            self.activation.initialise_layer(input_shape=self.output_shape, data_type=data_type, device=device)
+        if self.normalisation:
+            self.normalisation.initialise_layer(input_shape=self.output_shape, data_type=data_type, device=device)
+
+    def summary(self, offset=""):
+        if not hasattr(self, "input_shape"):
+            raise NotCompiledError("layer must be initialized correctly before calling layer.summary().")
+
+        input_shape = self.input_shape[0] if len(self.input_shape) == 1 else self.input_shape
+        output_shape = self.output_shape[0] if len(self.output_shape) == 1 else (self.output_shape if len(self.output_shape) != 0 else 1)
+        params_summary = " - Parameters: " + str(self.nparams) if self.nparams > 0 else ""
+        sublayer_offset = offset + "    "
+        normalisation_summary = ("\n" + self.normalisation.summary(sublayer_offset)) if self.normalisation else ""
+        activation_summary = ("\n" + self.activation.summary(sublayer_offset)) if self.activation else ""
+        return offset + f"{self.name} - (Input, Output): ({input_shape}, {output_shape})" + params_summary + normalisation_summary + activation_summary
+
+    def forward(self, input, **kwargs):
+        if not isinstance(input, torch.Tensor):
+            raise TypeError("input must be a torch.Tensor.")
+        if input.shape[1:] != self.input_shape:
+            raise ValueError(f"input is not the same shape as the spesified input_shape ({input.shape[1:], self.input_shape}).")
+
+        self.input = input
+        return self.input
+
+    def backward(self, dCdy, **kwargs):
+        if not isinstance(dCdy, torch.Tensor):
+            raise TypeError("dCdy must be a torch.Tensor.")
+        if dCdy.shape[1:] != self.output_shape and not (dCdy.ndim == 1 and self.output_shape[0] == 0):
+            raise ValueError(f"dCdy is not the same shape as the spesified output_shape ({dCdy.shape[1:], self.output_shape}).")
+
+        return dCdy
+    
+    def get_nparams(self):
+        if not hasattr(self, "nparams"):
+            raise NotCompiledError("layer must be initialized correctly before calling layer.summary().")
+
+        return self.nparams + (self.normalisation.nparams if self.normalisation else 0)
+
+    def get_parameters(self):
+        return tuple() if self.normalisation is None else self.normalisation.get_parameters()
+    
+    def clone(self):
+        return deepcopy(self)