src/transformers/models/trajectory_transformer/modeling_trajectory_transformer.py/CausalSelfAttention/forward
class CausalSelfAttention:    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