code: def __init__(self, config):
      super().__init__()
  
      if config.n_embd % config.n_head != 0:
          raise ValueError(
              f"n_head ({config.n_head}) should be a divisor of n_embd ({config.n_embd})"
          )
  
      # key, query, value projections for all heads
      self.key = nn.Linear(config.n_embd, config.n_embd)
      self.query = nn.Linear(config.n_embd, config.n_embd)
      self.value = nn.Linear(config.n_embd, config.n_embd)
  
      # regularization
      self.attn_drop = nn.Dropout(config.attn_pdrop)
      self.resid_drop = nn.Dropout(config.resid_pdrop)
  
      # output projection
      self.proj = nn.Linear(config.n_embd, config.n_embd)
  
      # causal mask to ensure that attention is only applied to the left in the input sequence
      self.register_buffer(
          "mask",
          torch.tril(
              torch.ones(config.block_size, config.block_size)
          ).view(1, 1, config.block_size, config.block_size),
      )
  
      # mask previous value estimates
      joined_dim = config.observation_dim + config.action_dim + 2
      self.mask.squeeze()[:, joined_dim - 1 :: joined_dim] = 0
  
      self.n_head = config.n_head
  
imports: import math
  import os
  from dataclasses import dataclass
  from typing import Optional, Tuple
  import numpy
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import functional
  from ...modeling_utils import PreTrainedModel
  from ...utils import ModelOutput, add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from .configuration_trajectory_transformer import TrajectoryTransformerConfig
label_desc: Use image node_id 1 to create a new CausalSelfAttention object from inherited base classes: nn with example: obj = CausalSelfAttention(config)