class_docstr: Multi Headed Attention Block in Transformers.
code: def forward(self, x):
      h = self.n_heads
      q, k, v = self.to_qkv(x).chunk(3, dim=-1)
      q, k, v = (
          rearrange(t, "b n (h d) -> b h n d", h=h) for t in (q, k, v)
      )
      sim = einsum("b h i d, b h j d -> b h i j", q, k) * self.scale
  
      attn = sim.softmax(dim=-1)
      attn = self.dropout(attn)
      if self.keep_attn:
          self.attn_weights = attn
      out = einsum("b h i j, b h j d -> b h i d", attn, v)
      out = rearrange(out, "b h n d -> b n (h d)", h=h)
      return self.to_out(out)
  
imports: import math
  from typing import Optional
  import torch
  from einops import rearrange
  from torch import einsum, nn
  from pytorch_tabular.utils import _initialize_kaiming
  from .gated_units import GEGLU, PositionWiseFeedForward, ReGLU, SwiGLU
label_desc: Use image node_id 2 for calling the MultiHeadedAttention obj's underlying member method code with example usage: obj.forward(x) and returns: self