class_docstr: A multi-head attention layer described in [1].
docstr: in_dim: Dimension of input
  out_dim: Dimension of output
  num_heads: Number of attention heads
  head_dim: Output dimension of each attention head
docstr_tok: ["in_dim",":","Dimension","of","input","out_dim",":","Dimension","of","output","num_heads",":","Number","of","attention","heads","head_dim",":","Output","dimension","of","each","attention","head"]
code: def __init__(
      self,
      in_dim: int,
      out_dim: int,
      num_heads: int,
      head_dim: int,
      **kwargs,
  ):
      """
      in_dim: Dimension of input
      out_dim: Dimension of output
      num_heads: Number of attention heads
      head_dim: Output dimension of each attention head
      """
      super().__init__(**kwargs)
  
      # No bias or non-linearity.
      self._num_heads = num_heads
      self._head_dim = head_dim
      self._qkv_layer = SlimFC(
          in_size=in_dim,
          out_size=3 * num_heads * head_dim,
          use_bias=False,
      )
  
      self._linear_layer = SlimFC(
          in_size=num_heads * head_dim, out_size=out_dim, use_bias=False
      )
  
file_path: ray/rllib/models/torch/modules/multi_head_attention.py
imports: from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.models.torch.misc import SlimFC
  from ray.rllib.utils.torch_utils import sequence_mask
  from ray.rllib.utils.framework import TensorType
label_desc: Use image node_id 1 to create a new MultiHeadAttention object from inherited base classes: nn with example: obj = MultiHeadAttention(in_dim, out_dim, num_heads, head_dim)