class_docstr: Multi-headed attention (MHA) module.
  
  This module is intended for attending over sequences of vectors.
  
  Rough sketch:
  - Compute keys (K), queries (Q), and values (V) as projections of inputs.
  - Attention weights are computed as W = softmax(QK^T / sqrt(key_size)).
  - Output is another projection of WV^T.
  
  For more detail, see the original Transformer paper:
    "Attention is all you need" https://arxiv.org/abs/1706.03762.
  
  Glossary of shapes:
  - T: Sequence length.
  - D: Vector (embedding) size.
  - H: Number of attention heads.
docstr: Initialises the module.
  
  Args:
    num_heads: Number of independent attention heads (H).
    key_size: The size of keys (K) and queries used for attention.
    w_init_scale: DEPRECATED. Please use w_init instead.
    w_init: Initialiser for weights in the linear map. Once `w_init_scale` is
      fully deprecated `w_init` will become mandatory. Until then it has a
      default value of `None` for backwards compatability.
    with_bias: Whether to add a bias when computing various linear
      projections.
    b_init: Optional initializer for bias. By default, zero.
    value_size: Optional size of the value projection (V). If None, defaults
      to the key size (K).
    model_size: Optional size of the output embedding (D'). If None, defaults
      to the key size multiplied by the number of heads (K * H).
    name: Optional name for this module.
docstr_tok: ["Initialises","the","module",".","Args",":","num_heads",":","Number","of","independent","attention","heads","(","H",")",".","key_size",":","The","size","of","keys","(","K",")","and","queries","used","for","attention",".","w_init_scale",":","DEPRECATED",".","Please","use","w_init","instead",".","w_init",":","Initialiser","for","weights","in","the","linear","map",".","Once","`","w_init_scale","`","is","fully","deprecated","`","w_init","`","will","become","mandatory",".","Until","then","it","has","a","default","value","of","`","None","`","for","backwards","compatability",".","with_bias",":","Whether","to","add","a","bias","when","computing","various","linear","projections",".","b_init",":","Optional","initializer","for","bias",".","By","default",",","zero",".","value_size",":","Optional","size","of","the","value","projection","(","V",")",".","If","None",",","defaults","to","the","key","size","(","K",")",".","model_size",":","Optional","size","of","the","output","embedding","(","D","'",")",".","If","None",",","defaults","to","the","key","size","multiplied","by","the","number","of","heads","(","K","*","H",")",".","name",":","Optional","name","for","this","module","."]
code: def __init__(
      self,
      num_heads: int,
      key_size: int,
      # TODO(b/240019186): Remove `w_init_scale`.
      w_init_scale: Optional[float] = None,
      *,
      w_init: Optional[hk.initializers.Initializer] = None,
      with_bias: bool = True,
      b_init: Optional[hk.initializers.Initializer] = None,
      value_size: Optional[int] = None,
      model_size: Optional[int] = None,
      name: Optional[str] = None,
  ):
      """Initialises the module.
  
      Args:
        num_heads: Number of independent attention heads (H).
        key_size: The size of keys (K) and queries used for attention.
        w_init_scale: DEPRECATED. Please use w_init instead.
        w_init: Initialiser for weights in the linear map. Once `w_init_scale` is
          fully deprecated `w_init` will become mandatory. Until then it has a
          default value of `None` for backwards compatability.
        with_bias: Whether to add a bias when computing various linear
          projections.
        b_init: Optional initializer for bias. By default, zero.
        value_size: Optional size of the value projection (V). If None, defaults
          to the key size (K).
        model_size: Optional size of the output embedding (D'). If None, defaults
          to the key size multiplied by the number of heads (K * H).
        name: Optional name for this module.
      """
      super().__init__(name=name)
      self.num_heads = num_heads
      self.key_size = key_size
      self.value_size = value_size or key_size
      self.model_size = model_size or key_size * num_heads
  
      # Backwards-compatibility for w_init_scale.
      if w_init_scale is not None:
          warnings.warn(
              "w_init_scale is deprecated; please pass an explicit weight "
              "initialiser instead.",
              DeprecationWarning,
          )
      if w_init and w_init_scale:
          raise ValueError(
              "Please provide only `w_init`, not `w_init_scale`."
          )
      if w_init is None and w_init_scale is None:
          raise ValueError(
              "Please provide a weight initializer: `w_init`. "
              "`w_init` will become mandatory once `w_init_scale` is "
              "fully deprecated."
          )
      if w_init is None:
          w_init = hk.initializers.VarianceScaling(w_init_scale)
      self.w_init = w_init
      self.with_bias = with_bias
      self.b_init = b_init
  
file_path: dm-haiku/haiku/_src/attention.py
imports: from typing import Optional
  import warnings
  from haiku._src import basic
  from haiku._src import initializers
  from haiku._src import module
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 to create a new MultiHeadAttention object from inherited base classes: hk with example: obj = MultiHeadAttention(num_heads, key_size, w_init_scale)