class_docstr: A RelativeMultiHeadAttention layer as described in [3].
  
  Uses segment level recurrence with state reuse.
docstr: Initializes a RelativeMultiHeadAttention nn.Module object.
  
  Args:
      in_dim (int):
      out_dim: The output dimension of this module. Also known as
          "attention dim".
      num_heads: The number of attention heads to use.
          Denoted `H` in [2].
      head_dim: The dimension of a single(!) attention head
          Denoted `D` in [2].
      input_layernorm: Whether to prepend a LayerNorm before
          everything else. Should be True for building a GTrXL.
      output_activation (Union[str, callable]): Optional activation
          function or activation function specifier (str).
          Should be "relu" for GTrXL.
      **kwargs:
docstr_tok: ["Initializes","a","RelativeMultiHeadAttention","nn.Module","object",".","Args",":","in_dim","(","int",")",":","out_dim",":","The","output","dimension","of","this","module",".","Also","known","as","``","attention","dim","''",".","num_heads",":","The","number","of","attention","heads","to","use",".","Denoted","`","H","`","in","[","2","]",".","head_dim",":","The","dimension","of","a","single","(","!",")","attention","head","Denoted","`","D","`","in","[","2","]",".","input_layernorm",":","Whether","to","prepend","a","LayerNorm","before","everything","else",".","Should","be","True","for","building","a","GTrXL",".","output_activation","(","Union","[","str",",","callable","]",")",":","Optional","activation","function","or","activation","function","specifier","(","str",")",".","Should","be","``","relu","''","for","GTrXL",".","*","*","kwargs",":"]
code: def __init__(
      self,
      in_dim: int,
      out_dim: int,
      num_heads: int,
      head_dim: int,
      input_layernorm: bool = False,
      output_activation: Union[str, callable] = None,
      **kwargs,
  ):
      """Initializes a RelativeMultiHeadAttention nn.Module object.
  
      Args:
          in_dim (int):
          out_dim: The output dimension of this module. Also known as
              "attention dim".
          num_heads: The number of attention heads to use.
              Denoted `H` in [2].
          head_dim: The dimension of a single(!) attention head
              Denoted `D` in [2].
          input_layernorm: Whether to prepend a LayerNorm before
              everything else. Should be True for building a GTrXL.
          output_activation (Union[str, callable]): Optional activation
              function or activation function specifier (str).
              Should be "relu" for GTrXL.
          **kwargs:
      """
      super().__init__(**kwargs)
  
      # No bias or non-linearity.
      self._num_heads = num_heads
      self._head_dim = head_dim
  
      # 3=Query, key, and value inputs.
      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,
          activation_fn=output_activation,
      )
  
      self._uvar = nn.Parameter(torch.zeros(num_heads, head_dim))
      self._vvar = nn.Parameter(torch.zeros(num_heads, head_dim))
      nn.init.xavier_uniform_(self._uvar)
      nn.init.xavier_uniform_(self._vvar)
      self.register_parameter("_uvar", self._uvar)
      self.register_parameter("_vvar", self._vvar)
  
      self._pos_proj = SlimFC(
          in_size=in_dim, out_size=num_heads * head_dim, use_bias=False
      )
      self._rel_pos_embedding = RelativePositionEmbedding(out_dim)
  
      self._input_layernorm = None
      if input_layernorm:
          self._input_layernorm = torch.nn.LayerNorm(in_dim)
  
file_path: ray/python/ray/rllib/models/torch/modules/relative_multi_head_attention.py
imports: from typing import Union
  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.typing import TensorType
label_desc: Use image node_id 1 to create a new RelativeMultiHeadAttention object from inherited base classes: nn with example: obj = RelativeMultiHeadAttention(in_dim, out_dim, num_heads, head_dim, input_layernorm, output_activation)