class_docstr: <a id="MHA"></a>
  
  ## Multi-Head Attention Module
  
  This computes scaled multi-headed attention for given `query`, `key` and `value` vectors.
  
  $$\mathop{Attention}(Q, K, V) = \underset{seq}{\mathop{softmax}}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V$$
  
  In simple terms, it finds keys that matches the query, and gets the values of
   those keys.
  
  It uses dot-product of query and key as the indicator of how matching they are.
  Before taking the $softmax$ the dot-products are scaled by $\frac{1}{\sqrt{d_k}}$.
  This is done to avoid large dot-product values causing softmax to
  give very small gradients when $d_k$ is large.
  
  Softmax is calculated along the axis of of the sequence (or time).
class_docstr_tok: ["<","a","id=","''","MHA","''",">","<","\/a",">","#","#","Multi-Head","Attention","Module","This","computes","scaled","multi-headed","attention","for","given","`","query","`",",","`","key","`","and","`","value","`","vectors",".","$","$","\\mathop","{","Attention","}","(","Q",",","K",",","V",")","=","\\underset","{","seq","}","{","\\mathop","{","softmax","}","}","\\Bigg","(","\\frac","{","Q","K^\\top","}","{","\\sqrt","{","d_k","}","}","\\Bigg",")","V","$","$","In","simple","terms",",","it","finds","keys","that","matches","the","query",",","and","gets","the","values","of","those","keys",".","It","uses","dot-product","of","query","and","key","as","the","indicator","of","how","matching","they","are",".","Before","taking","the","$","softmax","$","the","dot-products","are","scaled","by","$","\\frac","{","1","}","{","\\sqrt","{","d_k","}","}","$",".","This","is","done","to","avoid","large","dot-product","values","causing","softmax","to","give","very","small","gradients","when","$","d_k","$","is","large",".","Softmax","is","calculated","along","the","axis","of","of","the","sequence","(","or","time",")","."]
docstr: * `heads` is the number of heads.
  * `d_model` is the number of features in the `query`, `key` and `value` vectors.
docstr_tok: ["*","`","heads","`","is","the","number","of","heads",".","*","`","d_model","`","is","the","number","of","features","in","the","`","query","`",",","`","key","`","and","`","value","`","vectors","."]
code: def __init__(
      self,
      heads: int,
      d_model: int,
      dropout_prob: float = 0.1,
      bias: bool = True,
  ):
      """
      * `heads` is the number of heads.
      * `d_model` is the number of features in the `query`, `key` and `value` vectors.
      """
  
      super().__init__()
  
      # Number of features per head
      self.d_k = d_model // heads
      # Number of heads
      self.heads = heads
  
      # These transform the `query`, `key` and `value` vectors for multi-headed attention.
      self.query = PrepareForMultiHeadAttention(
          d_model, heads, self.d_k, bias=bias
      )
      self.key = PrepareForMultiHeadAttention(
          d_model, heads, self.d_k, bias=bias
      )
      self.value = PrepareForMultiHeadAttention(
          d_model, heads, self.d_k, bias=True
      )
  
      # Softmax for attention along the time dimension of `key`
      self.softmax = nn.Softmax(dim=1)
  
      # Output layer
      self.output = nn.Linear(d_model, d_model)
      # Dropout
      self.dropout = nn.Dropout(dropout_prob)
      # Scaling factor before the softmax
      self.scale = 1 / math.sqrt(self.d_k)
  
      # We store attentions so that it can be used for logging, or other computations if needed
      self.attn = None
  
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 MultiHeadAttention object from inherited base classes: nn with example: obj = MultiHeadAttention(heads, d_model, dropout_prob, bias)