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 forward(self, x: torch.Tensor):
      # Input has shape `[seq_len, batch_size, d_model]` or `[batch_size, d_model]`.
      # We apply the linear transformation to the last dimension and split that into
      # the heads.
      head_shape = x.shape[:-1]
  
      # Linear transform
      x = self.linear(x)
  
      # Split last dimension into heads
      x = x.view(*head_shape, self.heads, self.d_k)
  
      # Output has shape `[seq_len, batch_size, heads, d_k]` or `[batch_size, heads, d_model]`
      return 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 2 for calling the PrepareForMultiHeadAttention obj's underlying member method code with example usage: obj.forward(x) and returns: x