class_docstr: Multihead Attention (self-/cross-) with optional 'linear' attention.
  
  To learn more about Multihead Attention, see [1]. See the implementation
  of `Transformer` and the examples below to learn how to use the compression technique
  from [2] to speed up the module when the number of tokens is large.
  
  References:
  ----------
  [1] Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding" 2018
  [2] Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, Hao Ma "Linformer: Self-Attention with Linear Complexity", 2020
class_docstr_tok: ["Multihead","Attention","(","self-\/cross-",")","with","optional","'linear","'","attention",".","To","learn","more","about","Multihead","Attention",",","see","[","1","]",".","See","the","implementation","of","`","Transformer","`","and","the","examples","below","to","learn","how","to","use","the","compression","technique","from","[","2","]","to","speed","up","the","module","when","the","number","of","tokens","is","large",".","References",":","--","--","--","--","--","[","1","]","Jacob","Devlin",",","Ming-Wei","Chang",",","Kenton","Lee",",","Kristina","Toutanova","``","BERT",":","Pre-training","of","Deep","Bidirectional","Transformers","for","Language","Understanding","''","2018","[","2","]","Sinong","Wang",",","Belinda","Z.","Li",",","Madian","Khabsa",",","Han","Fang",",","Hao","Ma","``","Linformer",":","Self-Attention","with","Linear","Complexity","''",",","2020"]
docstr: Perform the forward pass.
  
  Parameters
  ----------
  x_q:
      query tokens
  x_kv:
      key-value tokens
  key_compression:
      Linformer-style compression for keys
  value_compression:
      Linformer-style compression for values
  
  Returns:
  ----------
      (tokens, attention_stats)
code: def forward(
      self,
      x_q: Tensor,
      x_kv: Tensor,
      key_compression: Optional[nn.Linear],
      value_compression: Optional[nn.Linear],
  ) -> Tuple[Tensor, Dict[str, Tensor]]:
      """Perform the forward pass.
  
      Parameters
      ----------
      x_q:
          query tokens
      x_kv:
          key-value tokens
      key_compression:
          Linformer-style compression for keys
      value_compression:
          Linformer-style compression for values
  
      Returns:
      ----------
          (tokens, attention_stats)
      """
      assert _all_or_none(
          [key_compression, value_compression]
      ), "If key_compression is (not) None, then value_compression must (not) be None"
      q, k, v = self.W_q(x_q), self.W_k(x_kv), self.W_v(x_kv)
      for tensor in [q, k, v]:
          assert (
              tensor.shape[-1] % self.n_heads == 0
          ), _INTERNAL_ERROR_MESSAGE
      if key_compression is not None:
          k = key_compression(k.transpose(1, 2)).transpose(1, 2)
          v = value_compression(v.transpose(1, 2)).transpose(1, 2)  # type: ignore
  
      batch_size = len(q)
      d_head_key = k.shape[-1] // self.n_heads
      d_head_value = v.shape[-1] // self.n_heads
      n_q_tokens = q.shape[1]
  
      q = self._reshape(q)
      k = self._reshape(k)
      attention_logits = q @ k.transpose(1, 2) / math.sqrt(d_head_key)
      attention_probs = F.softmax(attention_logits, dim=-1)
      if self.dropout is not None:
          attention_probs = self.dropout(attention_probs)
      x = attention_probs @ self._reshape(v)
      x = (
          x.reshape(batch_size, self.n_heads, n_q_tokens, d_head_value)
          .transpose(1, 2)
          .reshape(batch_size, n_q_tokens, self.n_heads * d_head_value)
      )
      if self.W_out is not None:
          x = self.W_out(x)
      return x, {
          "attention_logits": attention_logits,
          "attention_probs": attention_probs,
      }
  
imports: import enum
  import math
  import warnings
  from typing import Callable, Dict, List, Optional, Tuple, Union, cast
  import torch
  import torch.nn.functional
  from torch import Tensor, nn
  from .utils import init_weights
label_desc: Use image node_id 3 for calling the MultiheadAttention obj's underlying member method code with example usage: obj.forward(x_q, x_kv, key_compression, value_compression) and returns: x,  dict