class_docstr: <a id="PrepareMHA"></a>
  
  ## Prepare for multi-head attention
  
  This module does a linear transformation and splits the vector into given
  number of heads for multi-head attention.
  This is used to transform **key**, **query**, and **value** vectors.
class_docstr_tok: ["<","a","id=","''","PrepareMHA","''",">","<","\/a",">","#","#","Prepare","for","multi-head","attention","This","module","does","a","linear","transformation","and","splits","the","vector","into","given","number","of","heads","for","multi-head","attention",".","This","is","used","to","transform","*","*","key","*","*",",","*","*","query","*","*",",","and","*","*","value","*","*","vectors","."]
code: def __init__(self, d_model: int, heads: int, d_k: int, bias: bool):
      super().__init__()
      # Linear layer for linear transform
      self.linear = nn.Linear(d_model, heads * d_k, bias=bias)
      # Number of heads
      self.heads = heads
      # Number of dimensions in vectors in each head
      self.d_k = d_k
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/mha.py
imports: import math
  from typing import Optional, List
  import torch
  from torch import nn
  from labml import tracker
label_desc: Use image node_id 1 to create a new PrepareForMultiHeadAttention object from inherited base classes: nn with example: obj = PrepareForMultiHeadAttention(d_model, heads, d_k, bias)