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.
code: def _linear_projection(
      self,
      x: jax.Array,
      head_size: int,
      name: Optional[str] = None,
  ) -> jax.Array:
      y = hk.Linear(
          self.num_heads * head_size,
          w_init=self.w_init,
          with_bias=self.with_bias,
          b_init=self.b_init,
          name=name,
      )(x)
      *leading_dims, _ = x.shape
      return y.reshape((*leading_dims, self.num_heads, head_size))
  
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 3 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj._linear_projection(x, head_size, name) and returns: y