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: `query`, `key` and `value` are the tensors that store
  collection of *query*, *key* and *value* vectors.
  They have shape `[seq_len, batch_size, d_model]`.
  
  `mask` has shape `[seq_len, seq_len, batch_size]` and
  `mask[i, j, b]` indicates whether for batch `b`,
  query at position `i` has access to key-value at position `j`.
docstr_tok: ["`","query","`",",","`","key","`","and","`","value","`","are","the","tensors","that","store","collection","of","*","query","*",",","*","key","*","and","*","value","*","vectors",".","They","have","shape","`","[","seq_len",",","batch_size",",","d_model","]","`",".","`","mask","`","has","shape","`","[","seq_len",",","seq_len",",","batch_size","]","`","and","`","mask","[","i",",","j",",","b","]","`","indicates","whether","for","batch","`","b","`",",","query","at","position","`","i","`","has","access","to","key-value","at","position","`","j","`","."]
code: def forward(
      self,
      *,
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      mask: Optional[torch.Tensor] = None,
  ):
      """
      `query`, `key` and `value` are the tensors that store
      collection of *query*, *key* and *value* vectors.
      They have shape `[seq_len, batch_size, d_model]`.
  
      `mask` has shape `[seq_len, seq_len, batch_size]` and
      `mask[i, j, b]` indicates whether for batch `b`,
      query at position `i` has access to key-value at position `j`.
      """
  
      # `query`, `key` and `value`  have shape `[seq_len, batch_size, d_model]`
      seq_len, batch_size, _ = query.shape
  
      if mask is not None:
          mask = self.prepare_mask(mask, query.shape, key.shape)
  
      # Prepare `query`, `key` and `value` for attention computation.
      # These will then have shape `[seq_len, batch_size, heads, d_k]`.
      query = self.query(query)
      key = self.key(key)
      value = self.value(value)
  
      # Compute attention scores $Q K^\top$.
      # This gives a tensor of shape `[seq_len, seq_len, batch_size, heads]`.
      scores = self.get_scores(query, key)
  
      # Scale scores $\frac{Q K^\top}{\sqrt{d_k}}$
      scores *= self.scale
  
      # Apply mask
      if mask is not None:
          scores = scores.masked_fill(mask == 0, float("-inf"))
  
      # $softmax$ attention along the key sequence dimension
      # $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$
      attn = self.softmax(scores)
  
      # Save attentions if debugging
      tracker.debug("attn", attn)
  
      # Apply dropout
      attn = self.dropout(attn)
  
      # Multiply by values
      # $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V$$
      x = torch.einsum("ijbh,jbhd->ibhd", attn, value)
  
      # Save attentions for any other calculations
      self.attn = attn.detach()
  
      # Concatenate multiple heads
      x = x.reshape(seq_len, batch_size, -1)
  
      # Output layer
      return self.output(x)
  
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 4 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward() and returns: self