text: import torch
  from torch import nn
  import torch.nn.functional as F
  from typing import Callable, Tuple
  
  from .weight_init import trunc_normal_
  
  
  class RelPosSelfAttention(nn.Module):
      """
      Relative Position Self Attention
      From: https://gist.github.com/ShoufaChen/ec7b70038a6fdb488da4b34355380569
      """
  
      def __init__(self, h: int, w: int, dim: int, relative=True, fold_heads=False):
          super(RelPosSelfAttention, self).__init__()
          self.relative = relative
          self.fold_heads = fold_heads
          self.rel_emb_w = nn.Parameter(torch.Tensor(2 * w - 1, dim))
          self.rel_emb_h = nn.Parameter(torch.Tensor(2 * h - 1, dim))
  
          nn.init.normal_(self.rel_emb_w, std=dim ** -0.5)
          nn.init.normal_(self.rel_emb_h, std=dim ** -0.5)
  
      def forward(
              self,
              q: torch.Tensor,
              k: torch.Tensor,
              v: torch.Tensor,
              attn_mask: torch.Tensor
      ) -> torch.Tensor:
          """2D self-attention with rel-pos. Add option to fold heads."""
          bs, heads, h, w, dim = q.shape
          q = q * (dim ** -0.5)  # scaled dot-product
          logits = torch.einsum('bnhwd,bnpqd->bnhwpq', q, k)
          if self.relative:
              logits += self.relative_logits(q)
          weights = torch.reshape(logits, [-1, heads, h, w, h * w])
          weights = weights + torch.where(
              attn_mask.view(-1, 1, 1, 1, h * w),
              torch.zeros_like(weights),
              torch.zeros_like(weights) + float('-inf')
          )
          weights = F.softmax(weights, dim=-1)
          weights = torch.reshape(weights, [-1, heads, h, w, h, w])
          attn_out = torch.einsum('bnhwpq,bnpqd->bhwnd', weights, v)
          if self.fold_heads:
              attn_out = torch.reshape(attn_out, [-1, h, w, heads * dim])
          return attn_out
  
      def relative_logits(self, q):
          # Relative logits in width dimension.
          rel_logits_w = self.relative_logits_1d(
              q,
              self.rel_emb_w,
              transpose_mask=[0, 1, 2, 4, 3, 5]
          )
          # Relative logits in height dimension
          rel_logits_h = self.relative_logits_1d(
              q.permute(0, 1, 3, 2, 4),
              self.rel_emb_h,
              transpose_mask=[0, 1, 4, 2, 5, 3]
          )
          return rel_logits_h + rel_logits_w
  
      def relative_logits_1d(self, q, rel_k, transpose_mask):
          bs, heads, h, w, dim = q.shape
          rel_logits = torch.einsum('bhxyd,md->bhxym', q, rel_k)
          rel_logits = torch.reshape(rel_logits, [-1, heads * h, w, 2 * w - 1])
          rel_logits = self.rel_to_abs(rel_logits)
          rel_logits = torch.reshape(rel_logits, [-1, heads, h, w, w])
          rel_logits = torch.unsqueeze(rel_logits, dim=3)
          rel_logits = rel_logits.repeat(1, 1, 1, h, 1, 1)
          rel_logits = rel_logits.permute(*transpose_mask)
          return rel_logits
  
      def rel_to_abs(self, x):
          """
          Converts relative indexing to absolute.
          Input: [bs, heads, length, 2*length - 1]
          Output: [bs, heads, length, length]
          """
          bs, heads, length, _ = x.shape
          col_pad = torch.zeros((bs, heads, length, 1), dtype=x.dtype, device=x.device)
          x = torch.cat([x, col_pad], dim=3)
          flat_x = torch.reshape(x, [bs, heads, -1])
          flat_pad = torch.zeros((bs, heads, length - 1), dtype=x.dtype, device=x.device)
          flat_x_padded = torch.cat([flat_x, flat_pad], dim=2)
          final_x = torch.reshape(
              flat_x_padded, [bs, heads, length + 1, 2 * length - 1])
          final_x = final_x[:, :, :length, length - 1:]
          return final_x
  
  
  class GroupPointWise(nn.Module):
      """"""
      def __init__(self, in_channels, heads=4, proj_factor=1, target_dimension=None):
          super(GroupPointWise, self).__init__()
          if target_dimension is not None:
              proj_channels = target_dimension // proj_factor
          else:
              proj_channels = in_channels // proj_factor
          self.w = nn.Parameter(
              torch.Tensor(in_channels, heads, proj_channels // heads)
          )
  
          nn.init.normal_(self.w, std=0.01)
  
      def forward(self, x):
          # dim order:  pytorch BCHW v.s. TensorFlow BHWC
          x = x.permute(0, 2, 3, 1)
          """
          b: batch size
          h, w : imput height, width
          c: input channels
          n: num head
          p: proj_channel // heads
          """
          out = torch.einsum('bhwc,cnp->bnhwp', x, self.w)
          return out
  
  
  class RPSA(nn.Module):
      """
      """
      def __init__(self, in_channels, heads, height, width, pos_enc_type='relative'):
          super(RPSA, self).__init__()
          self.q_proj = GroupPointWise(in_channels, heads, proj_factor=1)
          self.k_proj = GroupPointWise(in_channels, heads, proj_factor=1)
          self.v_proj = GroupPointWise(in_channels, heads, proj_factor=1)
  
          assert pos_enc_type in ['relative', 'absolute']
          if pos_enc_type == 'relative':
              self.self_attention = RelPosSelfAttention(height, width, in_channels // heads, fold_heads=True)
          else:
              raise NotImplementedError
  
      def forward(self, x: torch.Tensor, input_mask: torch.Tensor) -> torch.Tensor:
          q = self.q_proj(x)
          k = self.k_proj(x)
          v = self.v_proj(x)
  
          o = self.self_attention(q=q, k=k, v=v, attn_mask=input_mask).permute(0, 3, 1, 2)
          return o
  
  
  class GPSA(nn.Module):
      """
      Gated positional self-attention
      From: https://github.com/facebookresearch/convit/blob/main/convit.py
      """
      def __init__(self, dim, height, width, n_heads=4, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.,
                   locality_strength=1., use_local_init=True):
          super().__init__()
          self.num_heads = n_heads
          self.dim = dim
          self.height = height
          self.width = width
          head_dim = dim // n_heads
          self.scale = qk_scale or head_dim ** -0.5
  
          self.qk = nn.Linear(dim, dim * 2, bias=qkv_bias)
          self.v = nn.Linear(dim, dim, bias=qkv_bias)
  
          self.attn_drop = nn.Dropout(attn_drop)
          self.proj = nn.Linear(dim, dim)
          self.pos_proj = nn.Linear(3, n_heads)
          self.proj_drop = nn.Dropout(proj_drop)
          self.locality_strength = locality_strength
          self.gating_param = nn.Parameter(torch.ones(self.num_heads))
          self.apply(self._init_weights)
          if use_local_init:
              self.local_init(locality_strength=locality_strength)
          self.register_buffer("rel_indices", self.get_rel_indices())
  
      def _init_weights(self, m):
          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)
  
      def forward(self, x: torch.Tensor, input_mask: torch.Tensor) -> torch.Tensor:
          x = torch.flatten(x, start_dim=-2, end_dim=-1).permute(0, 2, 1)
          input_mask = torch.flatten(input_mask, start_dim=-2, end_dim=-1)
          B, N, C = x.shape
  
          attn = self.get_attention(x, input_mask)
          v = self.v(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
          x = (attn @ v).transpose(1, 2).reshape(B, N, C)
          x = self.proj(x)
          x = self.proj_drop(x)
  
          x = x.permute(0, 2, 1)
          x = x.view(B, C, self.height, self.width)
          return x
  
      def get_attention(self, x, input_mask):
          B, N, C = x.shape
          qk = self.qk(x).reshape(B, N, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
          q, k = qk[0], qk[1]
          pos_score = self.rel_indices.expand(B, -1, -1, -1)
          pos_score = self.pos_proj(pos_score).permute(0, 3, 1, 2)
          patch_score = (q @ k.transpose(-2, -1)) * self.scale
          patch_score = patch_score.softmax(dim=-1)
          pos_score = pos_score.softmax(dim=-1)
  
          gating = self.gating_param.view(1, -1, 1, 1)
          attn = (1. - torch.sigmoid(gating)) * patch_score + torch.sigmoid(gating) * pos_score
          attn /= attn.sum(dim=-1).unsqueeze(-1)
          attn = self.attn_drop(attn)
          attn = attn * input_mask.float().unsqueeze(-2)
          attn = attn / attn.sum(dim=-1, keepdim=True)
          return attn
  
      def get_attention_map(self, x, input_mask, return_map=False):
          attn_map = self.get_attention(x, input_mask).mean(0)  # average over batch
          distances = self.rel_indices.squeeze()[:, :, -1] ** .5
          dist = torch.einsum('nm,hnm->h', (distances, attn_map))
          dist /= distances.size(0)
          if return_map:
              return dist, attn_map
          else:
              return dist
  
      def local_init(self, locality_strength=1.):
          self.v.weight.data.copy_(torch.eye(self.dim))
          locality_distance = 1  # max(1,1/locality_strength**.5)
  
          kernel_size = int(self.num_heads ** .5)
          center = (kernel_size - 1) / 2 if kernel_size % 2 == 0 else kernel_size // 2
          for h1 in range(kernel_size):
              for h2 in range(kernel_size):
                  position = h1 + kernel_size * h2
                  self.pos_proj.weight.data[position, 2] = -1
                  self.pos_proj.weight.data[position, 1] = 2 * (h1 - center) * locality_distance
                  self.pos_proj.weight.data[position, 0] = 2 * (h2 - center) * locality_distance
          self.pos_proj.weight.data *= locality_strength
  
      def get_rel_indices(self):
          assert self.height == self.width
          img_size = self.height
          num_patches = self.height * self.width
          rel_indices = torch.zeros(1, num_patches, num_patches, 3)
          ind = torch.arange(img_size).view(1, -1) - torch.arange(img_size).view(-1, 1)
          indx = ind.repeat(img_size, img_size)
          indy = ind.repeat_interleave(img_size, dim=0).repeat_interleave(img_size, dim=1)
          indd = indx ** 2 + indy ** 2
          rel_indices[:, :, :, 2] = indd.unsqueeze(0)
          rel_indices[:, :, :, 1] = indy.unsqueeze(0)
          rel_indices[:, :, :, 0] = indx.unsqueeze(0)
          device = self.qk.weight.device
          return rel_indices.to(device)
  
  
  class ViTBlock(nn.Module):
      def __init__(
              self,
              in_channels: int,
              out_channels: int,
              height: int,
              width: int,
              mhsa_layer: nn.Module,
              normalize: bool = True,
              activation: Callable = nn.GELU
      ):
          super(ViTBlock, self).__init__()
  
          self.norm1 = nn.LayerNorm([in_channels, height, width]) if normalize else nn.Identity()
          self.mhsa = mhsa_layer
  
          self.norm2 = nn.LayerNorm([in_channels, height, width]) if normalize else nn.Identity()
          self.mlp = nn.Sequential(
              nn.Conv2d(
                  in_channels=in_channels,
                  out_channels=out_channels,
                  kernel_size=(1, 1)
              ),
              activation(),
              nn.Conv2d(
                  in_channels=in_channels,
                  out_channels=out_channels,
                  kernel_size=(1, 1)
              )
          )
  
      def forward(self, x: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
          x, input_mask = x
          identity = x
          x = self.mhsa(self.norm1(x), input_mask)
          x = x + identity
          return (self.mlp(self.norm2(x)) + x) * input_mask, input_mask
  