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: ### Calculate scores between queries and keys
  
  This method can be overridden for other variations like relative attention.
docstr_tok: ["#","#","#","Calculate","scores","between","queries","and","keys","This","method","can","be","overridden","for","other","variations","like","relative","attention","."]
code: def get_scores(self, query: torch.Tensor, key: torch.Tensor):
      """
      ### Calculate scores between queries and keys
  
      This method can be overridden for other variations like relative attention.
      """
  
      # Calculate $Q K^\top$ or $S_{ijbh} = \sum_d Q_{ibhd} K_{jbhd}$
      return torch.einsum("ibhd,jbhd->ijbh", query, key)
  
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 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.get_scores(query, key) and returns: torch