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: Computes (optionally masked) MHA with queries, keys & values.
  
  This module broadcasts over zero or more 'batch-like' leading dimensions.
  
  Args:
    query: Embeddings sequence used to compute queries; shape [..., T', D_q].
    key: Embeddings sequence used to compute keys; shape [..., T, D_k].
    value: Embeddings sequence used to compute values; shape [..., T, D_v].
    mask: Optional mask applied to attention weights; shape [..., H=1, T', T].
  
  Returns:
    A new sequence of embeddings, consisting of a projection of the
      attention-weighted value projections; shape [..., T', D'].
docstr_tok: ["Computes","(","optionally","masked",")","MHA","with","queries",",","keys","&","values",".","This","module","broadcasts","over","zero","or","more","'batch-like","'","leading","dimensions",".","Args",":","query",":","Embeddings","sequence","used","to","compute","queries",";","shape","[","...",",","T","'",",","D_q","]",".","key",":","Embeddings","sequence","used","to","compute","keys",";","shape","[","...",",","T",",","D_k","]",".","value",":","Embeddings","sequence","used","to","compute","values",";","shape","[","...",",","T",",","D_v","]",".","mask",":","Optional","mask","applied","to","attention","weights",";","shape","[","...",",","H=1",",","T","'",",","T","]",".","Returns",":","A","new","sequence","of","embeddings",",","consisting","of","a","projection","of","the","attention-weighted","value","projections",";","shape","[","...",",","T","'",",","D","'","]","."]
code: def __call__(
      self,
      query: jax.Array,
      key: jax.Array,
      value: jax.Array,
      mask: Optional[jax.Array] = None,
  ) -> jax.Array:
      """Computes (optionally masked) MHA with queries, keys & values.
  
      This module broadcasts over zero or more 'batch-like' leading dimensions.
  
      Args:
        query: Embeddings sequence used to compute queries; shape [..., T', D_q].
        key: Embeddings sequence used to compute keys; shape [..., T, D_k].
        value: Embeddings sequence used to compute values; shape [..., T, D_v].
        mask: Optional mask applied to attention weights; shape [..., H=1, T', T].
  
      Returns:
        A new sequence of embeddings, consisting of a projection of the
          attention-weighted value projections; shape [..., T', D'].
      """
  
      # In shape hints below, we suppress the leading dims [...] for brevity.
      # Hence e.g. [A, B] should be read in every case as [..., A, B].
      *leading_dims, sequence_length, _ = query.shape
      projection = self._linear_projection
  
      # Compute key/query/values (overload K/Q/V to denote the respective sizes).
      query_heads = projection(
          query, self.key_size, "query"
      )  # [T', H, Q=K]
      key_heads = projection(key, self.key_size, "key")  # [T, H, K]
      value_heads = projection(
          value, self.value_size, "value"
      )  # [T, H, V]
  
      # Compute attention weights.
      attn_logits = jnp.einsum(
          "...thd,...Thd->...htT", query_heads, key_heads
      )
      attn_logits = attn_logits / np.sqrt(self.key_size).astype(
          key.dtype
      )
      if mask is not None:
          if mask.ndim != attn_logits.ndim:
              raise ValueError(
                  f"Mask dimensionality {mask.ndim} must match logits dimensionality "
                  f"{attn_logits.ndim}."
              )
          attn_logits = jnp.where(mask, attn_logits, -1e30)
      attn_weights = jax.nn.softmax(attn_logits)  # [H, T', T]
  
      # Weight the values by the attention and flatten the head vectors.
      attn = jnp.einsum(
          "...htT,...Thd->...thd", attn_weights, value_heads
      )
      attn = jnp.reshape(
          attn, (*leading_dims, sequence_length, -1)
      )  # [T', H*V]
  
      # Apply another projection to get the final embeddings.
      final_projection = hk.Linear(
          self.model_size,
          w_init=self.w_init,
          with_bias=self.with_bias,
          b_init=self.b_init,
      )
      return final_projection(attn)
  
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 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.__call__(query, key, value, mask) and returns: final_projection