class_docstr: Multi-Head Attention
class_docstr_tok: ["Multi-Head","Attention"]
code: def __init__(
      self, num_heads: int, d_model: int, dropout: float = 0.1
  ) -> None:
      super().__init__()
      assert (
          d_model % num_heads == 0
      ), "d_model must be divisible by num_heads"
  
      self.d_k = d_model // num_heads
      self.num_heads = num_heads
  
      self.linear_layers = nn.ModuleList(
          [nn.Linear(d_model, d_model) for _ in range(3)]
      )
      self.output_linear = nn.Linear(d_model, d_model)
  
imports: import math
  from typing import Any, Callable, Optional, Tuple
  import torch
  from torch import nn
label_desc: Use image node_id 1 to create a new MultiHeadAttention object from inherited base classes: nn with example: obj = MultiHeadAttention(num_heads, d_model, dropout)