class_docstr: Multi Headed Attention Block in Transformers.
code: def __init__(
      self,
      input_dim: int,
      num_heads: int = 8,
      head_dim: int = 16,
      dropout: int = 0.1,
      keep_attn: bool = True,
  ):
      super().__init__()
      assert (
          input_dim % num_heads == 0
      ), "'input_dim' must be multiples of 'num_heads'"
      inner_dim = head_dim * num_heads
      self.n_heads = num_heads
      self.scale = head_dim**-0.5
      self.keep_attn = keep_attn
  
      self.to_qkv = nn.Linear(input_dim, inner_dim * 3, bias=False)
      self.to_out = nn.Linear(inner_dim, input_dim)
  
      self.dropout = nn.Dropout(dropout)
  
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 1 to create a new MultiHeadedAttention object from inherited base classes: nn with example: obj = MultiHeadedAttention(input_dim, num_heads, head_dim, dropout, keep_attn)