class_docstr: Multi-headed Self Attention module.
  
  Source modified from:
  https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
class_docstr_tok: ["Multi-headed","Self","Attention","module",".","Source","modified","from",":","https",":","\/\/github.com\/rwightman\/pytorch-image-models\/blob\/master\/timm\/models\/vision_transformer.py"]
code: def forward(self, x: torch.Tensor) -> torch.Tensor:
      shape = x.shape
      B, C, H, W = shape
      N = H * W
      if len(shape) == 4:
          x = torch.flatten(x, start_dim=2).transpose(
              -2, -1
          )  # (B, N, C)
      qkv = (
          self.qkv(x)
          .reshape(B, N, 3, self.num_heads, self.head_dim)
          .permute(2, 0, 3, 1, 4)
      )
      q, k, v = qkv.unbind(
          0
      )  # make torchscript happy (cannot use tensor as tuple)
  
      # trick here to make q@k.t more stable
      attn = (q * self.scale) @ k.transpose(-2, -1)
      attn = attn.softmax(dim=-1)
      attn = self.attn_drop(attn)
  
      x = (attn @ v).transpose(1, 2).reshape(B, N, C)
      x = self.proj(x)
      x = self.proj_drop(x)
      if len(shape) == 4:
          x = x.transpose(-2, -1).reshape(B, C, H, W)
  
      return x
  
imports: import os
  import copy
  from functools import partial
  from typing import List, Tuple, Optional, Union
  import torch
  import torch.nn
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.models.layers import DropPath, trunc_normal_
  from timm.models.registry import register_model
  from models.modules.mobileone import MobileOneBlock
  from models.modules.replknet import ReparamLargeKernelConv
label_desc: Use image node_id 2 for calling the MHSA obj's underlying member method code with example usage: obj.forward(x) and returns: x