class_docstr: ## Multi-head attention with rotary positional embeddings
  
  We override [multi-head attention from original transformer](../mha.html).
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`.
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)
  
      # Apply dropout
      attn = self.dropout(attn)
  
      # Rotate value embeddings before taking the weighted sum so that they contain positional information
      value = self.value_rotary_pe(value)
  
      # Multiply by values
      # $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V$$
      x = torch.einsum(
          "ijbh,jbhd->ibhd", attn, self.value_rotary_pe(value)
      )
  
      # Rotate in the opposite direction so that each embedding hold the relative positions
      x = self.value_reverse_rotary_pe(x)
  
      # 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/rope/value_pe/__init__.py
imports: from typing import Optional
  import torch
  from labml_nn.transformers.rope import RotaryPositionalEmbeddings, RotaryPEMultiHeadAttention
label_desc: Use image node_id 2 for calling the RotaryValuePEMultiHeadAttention obj's underlying member method code with example usage: obj.forward() and returns: self