class_docstr: A multi-head attention layer described in [1].
code: def forward(self, inputs: TensorType) -> TensorType:
      L = list(inputs.size())[1]  # length of segment
      H = self._num_heads  # number of attention heads
      D = self._head_dim  # attention head dimension
  
      qkv = self._qkv_layer(inputs)
  
      queries, keys, values = torch.chunk(input=qkv, chunks=3, dim=-1)
      queries = queries[:, -L:]  # only query based on the segment
  
      queries = torch.reshape(queries, [-1, L, H, D])
      keys = torch.reshape(keys, [-1, L, H, D])
      values = torch.reshape(values, [-1, L, H, D])
  
      score = torch.einsum("bihd,bjhd->bijh", queries, keys)
      score = score / D**0.5
  
      # causal mask of the same length as the sequence
      mask = sequence_mask(torch.arange(1, L + 1), dtype=score.dtype)
      mask = mask[None, :, :, None]
      mask = mask.float()
  
      masked_score = score * mask + 1e30 * (mask - 1.0)
      wmat = nn.functional.softmax(masked_score, dim=2)
  
      out = torch.einsum("bijh,bjhd->bihd", wmat, values)
      shape = list(out.size())[:2] + [H * D]
      #        temp = torch.cat(temp2, [H * D], dim=0)
      out = torch.reshape(out, shape)
      return self._linear_layer(out)
  
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 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(inputs) and returns: self