code: def forward(self, x: torch.Tensor) -> torch.Tensor:
      B, N, C = x.shape
      # split keys, queries and values in num_heads
      # NOTE: the line below differs from timm
      # timm: qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
      qkv = (
          self.qkv(x)
          .reshape(B, N, 3, self.num_heads, C // self.num_heads)
          .permute(2, 0, 3, 1, 4)
      )
      q, k, v = qkv[0], qkv[1], qkv[2]
  
      # sum up over the last axis
      att = torch.einsum("bhqd, bhkd -> bhqk", q, k) * self.scale
      att = att.softmax(dim=-1)
      att = self.att_drop(att)
  
      # sum up over the third axis
      out = torch.einsum("bhal, bhlv -> bhav ", att, v)
      out = out.permute(0, 2, 1, 3).contiguous().view(B, N, -1)
      out = self.projection(out)
      out = self.projection_drop(out)
      return out
  
imports: from typing import Any, Callable, Dict, List, Optional, Tuple
  import torch
  from torch import nn
  from kornia.core import Module, Tensor
label_desc: Use image node_id 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(x) and returns: out