class_docstr: A RelativeMultiHeadAttention layer as described in [3].
  
  Uses segment level recurrence with state reuse.
class_docstr_tok: ["A","RelativeMultiHeadAttention","layer","as","described","in","[","3","]",".","Uses","segment","level","recurrence","with","state","reuse","."]
code: def forward(
      self, inputs: TensorType, memory: TensorType = None
  ) -> TensorType:
      T = list(inputs.size())[1]  # length of segment (time)
      H = self._num_heads  # number of attention heads
      d = self._head_dim  # attention head dimension
  
      # Add previous memory chunk (as const, w/o gradient) to input.
      # Tau (number of (prev) time slices in each memory chunk).
      Tau = list(memory.shape)[1]
      inputs = torch.cat((memory.detach(), inputs), dim=1)
  
      # Apply the Layer-Norm.
      if self._input_layernorm is not None:
          inputs = self._input_layernorm(inputs)
  
      qkv = self._qkv_layer(inputs)
  
      queries, keys, values = torch.chunk(input=qkv, chunks=3, dim=-1)
      # Cut out Tau memory timesteps from query.
      queries = queries[:, -T:]
  
      queries = torch.reshape(queries, [-1, T, H, d])
      keys = torch.reshape(keys, [-1, Tau + T, H, d])
      values = torch.reshape(values, [-1, Tau + T, H, d])
  
      R = self._pos_proj(self._rel_pos_embedding(Tau + T))
      R = torch.reshape(R, [Tau + T, H, d])
  
      # b=batch
      # i and j=time indices (i=max-timesteps (inputs); j=Tau memory space)
      # h=head
      # d=head-dim (over which we will reduce-sum)
      score = torch.einsum(
          "bihd,bjhd->bijh", queries + self._uvar, keys
      )
      pos_score = torch.einsum(
          "bihd,jhd->bijh", queries + self._vvar, R
      )
      score = score + self.rel_shift(pos_score)
      score = score / d**0.5
  
      # causal mask of the same length as the sequence
      mask = sequence_mask(
          torch.arange(Tau + 1, Tau + T + 1), dtype=score.dtype
      ).to(score.device)
      mask = mask[None, :, :, None]
  
      masked_score = score * mask + 1e30 * (mask.float() - 1.0)
      wmat = nn.functional.softmax(masked_score, dim=2)
  
      out = torch.einsum("bijh,bjhd->bihd", wmat, values)
      shape = list(out.shape)[:2] + [H * d]
      out = torch.reshape(out, shape)
  
      return self._linear_layer(out)
  
file_path: 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 2 for calling the RelativeMultiHeadAttention obj's underlying member method code with example usage: obj.forward(inputs, memory) and returns: self