--- a
+++ b/biovil_t/transformer.py
@@ -0,0 +1,266 @@
+#  -------------------------------------------------------------------------------------------
+#  Copyright (c) Microsoft Corporation. All rights reserved.
+#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+#  -------------------------------------------------------------------------------------------
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Callable, Optional, Set, Tuple
+
+import torch
+import torch.nn as nn
+from timm.models.layers import DropPath, Mlp, trunc_normal_
+
+
+def torch_int_div(tensor1, tensor2):
+    """
+    A function that performs integer division across different versions of PyTorch.
+    """
+    return torch.div(tensor1, tensor2, rounding_mode="floor")
+
+@dataclass
+class MultiHeadAttentionOutput:
+    mha_output: torch.Tensor
+    attention: Optional[torch.Tensor] = None
+
+
+class VisionTransformerPooler(nn.Module):
+    """
+    :param input_dim: Input feature dimension (i.e., channels in old CNN terminology)
+    :param grid_shape: Shape of the grid of patches per image
+    :param num_heads: Number of self-attention heads within the MHA block
+    :param num_blocks: Number of blocks per attention layer
+    :param norm_layer: Normalisation layer
+
+    `self.type_embed`: Is used to characterise prior and current scans, and
+                       create permutation variance across modalities/series.
+    """
+
+    def __init__(self,
+                 input_dim: int,
+                 grid_shape: Tuple[int, int],
+                 num_heads: int = 8,
+                 num_blocks: int = 3,
+                 norm_layer: Any = partial(nn.LayerNorm, eps=1e-6)):
+        super().__init__()
+
+        block_kwargs = dict(dim=input_dim, num_heads=num_heads, mlp_ratio=1., drop=0.10, attn_drop=0.10,
+                            drop_path=0.25, act_layer=nn.GELU, norm_layer=norm_layer)
+        self.blocks = nn.ModuleList([Block(**block_kwargs) for _ in range(num_blocks)])
+        self.norm_post = norm_layer(input_dim)
+        self.grid_shape = grid_shape
+        self.num_patches = grid_shape[0] * grid_shape[1]
+        self.num_blocks = num_blocks
+
+        # Temporal positional embeddings
+        num_series: int = 2
+        self.type_embed = nn.Parameter(torch.zeros(num_series, 1, input_dim))
+        trunc_normal_(self.type_embed, std=.02)
+
+        # Positional embeddings 1 x L x C (L: Sequence length, C: Feature dimension)
+        self.pos_drop = nn.Dropout(p=0.10)
+        pos_embed_class = SinePositionEmbedding(embedding_dim=input_dim // 2, normalize=True)
+        pos_embed = pos_embed_class(mask=torch.ones([1, grid_shape[0], grid_shape[1]]))  # 1 x L x C
+        self.register_buffer("pos_embed", pos_embed, persistent=False)
+
+        # Initialisation
+        self.apply(self._init_weights)
+
+    def no_weight_decay(self) -> Set[str]:
+        return {'type_embed'}
+
+    def forward(self, current_image: torch.Tensor, previous_image: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, C, H, W = current_image.shape
+        assert H == self.grid_shape[0] and W == self.grid_shape[1], "Input and grid shapes do not match"
+
+        # Flatten patch embeddings to have shape (B x L x C), L = H * W
+        if previous_image is not None:
+            assert previous_image.shape == current_image.shape, "current_image and previous_image shapes do not match"
+            previous_image = previous_image.view(B, C, H * W).transpose(1, 2)
+        current_image = current_image.view(B, C, H * W).transpose(1, 2)
+        pos_embed = self.pos_embed.repeat(B, 1, 1)  # type: ignore
+
+        # Final token activations (B x 2L x C)
+        token_features = self.forward_after_reshape(x=current_image, pos_embed=pos_embed, x_previous=previous_image)
+
+        # Extract the patch features of current image
+        cur_img_token_id = 0
+        current_token_features = token_features[:, cur_img_token_id:self.num_patches+cur_img_token_id]
+        current_patch_features = current_token_features.transpose(1, 2).view(B, C, H, W)
+
+        return current_patch_features
+
+    def forward_after_reshape(self,
+                              x: torch.Tensor,
+                              pos_embed: torch.Tensor,
+                              x_previous: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, L, _ = x.shape  # Batch, Sequence length, Feature dimension
+
+        # Positional and type embeddings
+        type_embed = self.type_embed[0].expand(B, L, -1)
+        if x_previous is not None:
+            x = torch.cat((x, x_previous), dim=1)
+            pos_embed = torch.cat((pos_embed, pos_embed), dim=1)
+            prev_type_embed = self.type_embed[1].expand(B, L, -1)
+            type_embed = torch.cat((type_embed, prev_type_embed), dim=1)
+
+        # Add positional and type embeddings (used in query and key matching)
+        pos_and_type_embed = pos_embed + type_embed
+
+        # Positional dropout
+        x = self.pos_drop(x)
+
+        # Multihead attention followed by MLP
+        for block in self.blocks:
+            x = block(x=x, pos_and_type_embed=pos_and_type_embed)
+        x = self.norm_post(x)
+
+        return x
+
+    def _init_weights(self, m: nn.Module) -> None:
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+
+class MultiHeadAttentionLayer(nn.Module):
+    """
+    Multi-head self attention module
+
+    The content builds on top of the TIMM library (vision_transformer.py) and differs by the following:
+        - Defines a custom `MultiHeadAttentionLayer` which does not only apply `self-attention` but it can be
+            generalised to arbitrary (query, key, value) input tuples. This feature can be valuable to process
+            more than 2 scans at a time.
+        - `Self-attention` specific use-case can still be invoked by calling the `forward_as_mhsa` method.
+    """
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int = 8,
+                 qkv_bias: bool = False,
+                 attn_drop: float = 0.,
+                 proj_drop: float = 0.) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        assert dim % num_heads == 0, f"The embedding dim ({dim}) must be divisible by the number of heads ({num_heads})"
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+        self.return_attention = False
+
+        self.proj_q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.proj_k = nn.Linear(dim, dim, bias=qkv_bias)
+        self.proj_v = nn.Linear(dim, dim, bias=qkv_bias)
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, k: torch.Tensor, q: torch.Tensor, v: torch.Tensor) -> MultiHeadAttentionOutput:
+        B, N, C = v.shape
+        assert C % self.num_heads == 0, \
+            f"The embedding dim ({C}) must be divisible by the number of heads ({self.num_heads})"
+
+        w_q = self.proj_q(q).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        w_k = self.proj_k(k).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        w_v = self.proj_v(v).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+
+        attn = (w_q @ w_k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        o = (attn @ w_v).transpose(1, 2).reshape(B, N, C)
+        o = self.proj(o)
+        o = self.proj_drop(o)
+
+        attention_output = attn if self.return_attention else None
+
+        return MultiHeadAttentionOutput(mha_output=o, attention=attention_output)
+
+    def forward_as_mhsa(self, input: torch.Tensor) -> MultiHeadAttentionOutput:
+        return self(k=input, q=input, v=input)
+
+
+class Block(nn.Module):
+    """
+    Encapsulates multi-layer perceptron and multi-head self attention modules into a block.
+
+    The content builds on top of the TIMM library (vision_transformer.py) and differs by the following:
+        - This implementation uses spatio-temporal positional embeddings instead of 2D positional embeddings only,
+            and they are taken into account within the forward pass of each ViT block.
+        - Utilises the custom defined `MultiHeadAttentionLayer` which does not apply `self-attention` only but can be
+            generalised to arbitrary (query, key, value) tuples. This can be valuable to process more than 2 scans.
+
+    Positional and type embeddings are handled in a similar fashion as DETR object localisation paper
+    https://alcinos.github.io/detr_page/, where a fixed set of sine/cos positional embeddings are used
+    in an additive manner to Q and K tensors.
+    """
+
+    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 1., qkv_bias: bool = False, drop: float = 0.,
+                 attn_drop: float = 0., drop_path: float = 0., act_layer: Callable = nn.GELU,
+                 norm_layer: Callable = nn.LayerNorm) -> None:
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = MultiHeadAttentionLayer(dim=dim, num_heads=num_heads, qkv_bias=qkv_bias,
+                                            attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def with_pos_and_type_embed(self, tensor: torch.Tensor, emb: Optional[torch.Tensor]) -> torch.Tensor:
+        # Add positional embeddings to key and query tensors
+        return tensor if emb is None else tensor + emb
+
+    def forward(self, x: torch.Tensor, pos_and_type_embed: Optional[torch.Tensor]) -> torch.Tensor:
+        x_with_emb = self.with_pos_and_type_embed(self.norm1(x), emb=pos_and_type_embed)
+        x = x + self.drop_path(self.attn.forward_as_mhsa(x_with_emb).mha_output)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+
+class SinePositionEmbedding():
+    """
+    This is a more standard version of the position embedding, very similar to the one used by the Attention is all you
+    need paper, generalized to work on images.
+    """
+
+    def __init__(self,
+                 embedding_dim: int = 64,
+                 temperature: int = 10000,
+                 normalize: bool = False,
+                 scale: float = None) -> None:
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def __call__(self, mask: torch.Tensor) -> torch.Tensor:
+        assert mask is not None, "No pixel mask provided"
+        B, H, W = mask.shape
+        y_embed = mask.cumsum(1, dtype=torch.float32)
+        x_embed = mask.cumsum(2, dtype=torch.float32)
+        if self.normalize:
+            y_embed = y_embed / (y_embed[:, -1:, :] + 1e-6) * self.scale
+            x_embed = x_embed / (x_embed[:, :, -1:] + 1e-6) * self.scale
+
+        dim_t = torch.arange(self.embedding_dim, dtype=torch.float32)
+        dim_t = self.temperature ** (2 * torch_int_div(dim_t, 2) / self.embedding_dim)
+
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos = torch.cat((pos_y, pos_x), dim=3).view(B, H * W, self.embedding_dim * 2)
+
+        return pos