text: import torch
  import torch.nn as nn
  from einops import rearrange
  
  from self_attention_cv import TransformerEncoder
  from ..common import expand_to_batch
  
  
  class ViT(nn.Module):
      def __init__(self, *,
                   img_dim,
                   in_channels=3,
                   patch_dim=16,
                   num_classes=10,
                   dim=512,
                   blocks=6,
                   heads=4,
                   dim_linear_block=1024,
                   dim_head=None,
                   dropout=0, transformer=None, classification=True):
          """
          Minimal re-implementation of ViT
          Args:
              img_dim: the spatial image size
              in_channels: number of img channels
              patch_dim: desired patch dim
              num_classes: classification task classes
              dim: the linear layer's dim to project the patches for MHSA
              blocks: number of transformer blocks
              heads: number of heads
              dim_linear_block: inner dim of the transformer linear block
              dim_head: dim head in case you want to define it. defaults to dim/heads
              dropout: for pos emb and transformer
              transformer: in case you want to provide another transformer implementation
              classification: creates an extra CLS token that we will index in the final classification layer
          """
          super().__init__()
          assert img_dim % patch_dim == 0, f'patch size {patch_dim} not divisible by img dim {img_dim}'
          self.p = patch_dim
          self.classification = classification
          # tokens = number of patches
          tokens = (img_dim // patch_dim) ** 2
          self.token_dim = in_channels * (patch_dim ** 2)
          self.dim = dim
          self.dim_head = (int(self.dim / heads)) if dim_head is None else dim_head
  
          # Projection and pos embeddings
          self.project_patches = nn.Linear(self.token_dim, self.dim)
  
          self.emb_dropout = nn.Dropout(dropout)
  
          self.cls_token = nn.Parameter(torch.randn(1, 1, self.dim))
          self.pos_emb1D = nn.Parameter(torch.randn(tokens + 1, self.dim))
  
          if self.classification:
              self.mlp_head = nn.Linear(self.dim, num_classes)
  
          if transformer is None:
              self.transformer = TransformerEncoder(self.dim, blocks=blocks, heads=heads,
                                                    dim_head=self.dim_head,
                                                    dim_linear_block=dim_linear_block,
                                                    dropout=dropout)
          else:
              self.transformer = transformer
  
      def forward(self, img, mask=None):
          # Create patches
          # from [batch, channels, h, w] to [batch, tokens , N], N=p*p*c , tokens = h/p *w/p
          img_patches = rearrange(img,
                                  'b c (patch_x x) (patch_y y) -> b (x y) (patch_x patch_y c)',
                                  patch_x=self.p, patch_y=self.p)
  
          batch_size, tokens, _ = img_patches.shape
  
          # project patches with linear layer + add pos emb
          img_patches = self.project_patches(img_patches)
  
          img_patches = torch.cat((expand_to_batch(self.cls_token, desired_size=batch_size), img_patches), dim=1)
  
          # add pos. embeddings. + dropout
          # indexing with the current batch's token length to support variable sequences
          img_patches = img_patches + self.pos_emb1D[:tokens + 1, :]
          patch_embeddings = self.emb_dropout(img_patches)
  
          # feed patch_embeddings and output of transformer. shape: [batch, tokens, dim]
          y = self.transformer(patch_embeddings, mask)
  
          # we index only the cls token for classification. nlp tricks :P
          return self.mlp_head(y[:, 0, :]) if self.classification else y[:, 1:, :]
  