code: def forward(
      self,
      hidden_states: Optional[Tuple[torch.FloatTensor]],
      layer_past: Optional[Tuple[torch.Tensor]] = None,
      use_cache: Optional[bool] = False,
      output_attentions: Optional[bool] = False,
  ):
      batch_size, sequence_length, embedding_dim = hidden_states.size()
  
      # calculate query, key, values for all heads in batch and move head forward to be the batch dim
      # [ batch_size x n_heads x sequence_length x head_dim ]
      key = (
          self.key(hidden_states)
          .view(
              batch_size,
              sequence_length,
              self.n_head,
              embedding_dim // self.n_head,
          )
          .transpose(1, 2)
      )
      query = (
          self.query(hidden_states)
          .view(
              batch_size,
              sequence_length,
              self.n_head,
              embedding_dim // self.n_head,
          )
          .transpose(1, 2)
      )
      value = (
          self.value(hidden_states)
          .view(
              batch_size,
              sequence_length,
              self.n_head,
              embedding_dim // self.n_head,
          )
          .transpose(1, 2)
      )
  
      if layer_past is not None:
          past_key, past_value = layer_past
          key = torch.cat((past_key, key), dim=-2)
          value = torch.cat((past_value, value), dim=-2)
  
      if use_cache is True:
          present = (key, value)
      else:
          present = None
  
      # causal self-attention
      # [ batch_size x n_heads x sequence_length x sequence_length ]
      attn_weights = (torch.matmul(query, key.transpose(-2, -1))) * (
          1.0 / math.sqrt(key.size(-1))
      )
      attn_weights = attn_weights.masked_fill(
          self.mask[:, :, :sequence_length, :sequence_length] == 0,
          torch.finfo(attn_weights.dtype).min,
      )
      attn_weights = F.softmax(attn_weights, dim=-1)
      self._attn_map = attn_weights.clone()
      attn_weights = self.attn_drop(attn_weights)
  
      output = torch.matmul(attn_weights, value)
      # [ batch_size x sequence_length x embedding_dim ]
      # re-assemble all head outputs side by side
      output = (
          output.transpose(1, 2)
          .contiguous()
          .view(batch_size, sequence_length, embedding_dim)
      )
  
      # output projection
      output = self.resid_drop(self.proj(output))
  
      outputs = (output, present)
      if output_attentions:
          outputs += (attn_weights,)
  
      return outputs
  
imports: import math
  import os
  from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  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 2 for calling the CausalSelfAttention obj's underlying member method code with example usage: obj.forward(hidden_states, layer_past, use_cache, output_attentions) and returns: outputs