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 __init__(self, config: GPTConfig):
      super().__init__()
      assert config.n_embed % config.n_head == 0
      # key, query, value projections for all heads, but in a batch
      self.c_attn = nn.Linear(config.n_embed, 3 * config.n_embed)
      # output projection
      self.c_proj = nn.Linear(config.n_embed, config.n_embed)
      # regularization
      self.attn_dropout = nn.Dropout(config.attn_pdrop)
      self.resid_dropout = nn.Dropout(config.resid_pdrop)
      # causal mask to ensure that attention is only applied to the left
      # in the input sequence
      self.register_buffer(
          "bias",
          torch.tril(
              torch.ones(config.block_size, config.block_size)
          ).view(1, 1, config.block_size, config.block_size),
      )
      self.n_head = config.n_head
      self.n_embed = config.n_embed
  
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 1 to create a new CausalSelfAttention object from inherited base classes: nn with example: obj = CausalSelfAttention(config)