class_docstr: Vanilla multi-head masked self-attention layer with a projection at the end.
  It is possible to use torch.nn.MultiheadAttention here but I am including an
  explicit implementation here to show that there is nothing too scary here.
class_docstr_tok: ["Vanilla","multi-head","masked","self-attention","layer","with","a","projection","at","the","end",".","It","is","possible","to","use","torch.nn.MultiheadAttention","here","but","I","am","including","an","explicit","implementation","here","to","show","that","there","is","nothing","too","scary","here","."]
code: def forward(self, x, attention_masks=None):
      # batch size, sequence length, embedding dimensionality (n_embed)
      B, T, C = x.size()
  
      # calculate query, key, values for all heads in batch and move head
      # forward to be the batch dim
      q, k, v = self.c_attn(x).split(self.n_embed, dim=2)
      # (B, nh, T, hs)
      k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
      # (B, nh, T, hs)
      q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
      # (B, nh, T, hs)
      v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
  
      # causal self-attention; Self-attend:
      # (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
      att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
      att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float("-inf"))
      if attention_masks is not None:
          att = att + attention_masks
      att = F.softmax(att, dim=-1)
      att = self.attn_dropout(att)
      y = att @ v  # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
      # re-assemble all head outputs side by side
      y = y.transpose(1, 2).contiguous().view(B, T, C)
  
      # output projection
      y = self.resid_dropout(self.c_proj(y))
      return y, att
  
imports: import math
  from dataclasses import dataclass
  from typing import Tuple
  import torch
  import torch.nn
  from torch.nn import functional
  from ray.rllib.utils.annotations import DeveloperAPI
  from ray.rllib.utils.deprecation import Deprecated
label_desc: Use image node_id 2 for calling the CausalSelfAttention obj's underlying member method code with example usage: obj.forward(x, attention_masks) and returns: y,  att