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"]
code: def _reshape(self, x: Tensor) -> Tensor:
      batch_size, n_tokens, d = x.shape
      d_head = d // self.n_heads
      return (
          x.reshape(batch_size, n_tokens, self.n_heads, d_head)
          .transpose(1, 2)
          .reshape(batch_size * self.n_heads, n_tokens, d_head)
      )
  
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 2 for calling the MultiheadAttention obj's underlying member method code with example usage: obj._reshape(x) and returns: x