class_docstr: Multi-head self attention.
  Using torch or triton attention implemetation enables user to also use additive bias.
class_docstr_tok: ["Multi-head","self","attention",".","Using","torch","or","triton","attention","implemetation","enables","user","to","also","use","additive","bias","."]
code: def forward(
      self,
      hidden_states: torch.Tensor,
      position_bias: torch.Tensor,
      past_key_value: Optional[Tuple[torch.Tensor]] = None,
      attention_mask: Optional[torch.Tensor] = None,
  ):
      batch_size, seq_length = hidden_states.shape[:2]
  
      mixed_qkv = self.Wqkv(hidden_states)
      query_states, key_states, value_states = mixed_qkv.chunk(3, dim=2)
      query_states = query_states.reshape(
          batch_size, seq_length, self.n_heads, self.head_dim
      ).transpose(1, 2)
      key_states = key_states.reshape(
          batch_size, seq_length, self.n_heads, self.head_dim
      ).transpose(1, 2)
      value_states = value_states.reshape(
          batch_size, seq_length, self.n_heads, self.head_dim
      ).transpose(1, 2)
  
      if past_key_value is not None:
          if len(past_key_value) != 0:
              key_states = torch.cat(
                  [past_key_value[0], key_states], dim=2
              )
              value_states = torch.cat(
                  [past_key_value[1], value_states], dim=2
              )
          past_key_value = (key_states, value_states)
      else:
          past_key_value = (key_states, value_states)
  
      attention_scores = (
          torch.matmul(query_states, key_states.transpose(-1, -2))
          * self.softmax_scale
      )
  
      query_length = (
          seq_length
          if past_key_value is None
          else seq_length + past_key_value[0].shape[2]
      )
  
      if position_bias is not None:
          if len(position_bias.shape) != 3:
              raise ValueError(
                  f"Expecting position_bias shape to be 3 dimensions, got {len(position_bias.shape)}"
              )
          key_length = key_states.shape[-2]
  
          position_bias_query_index = max(
              0, position_bias.size(1) - query_length
          )
          position_bias_key_index = max(
              0, position_bias.size(2) - key_length
          )
  
          position_bias = position_bias[
              :, position_bias_query_index:, position_bias_key_index:
          ]
  
          attention_scores = attention_scores + position_bias
  
      if attention_mask is not None:
          attention_scores = attention_scores.masked_fill(
              attention_mask, torch.finfo(query_states.dtype).min
          )
  
      # (batch_size, n_heads, seq_length, key_length)
      attn_weights = nn.functional.softmax(
          attention_scores.float(), dim=-1
      ).to(value_states.dtype)
      attn_weights = nn.functional.dropout(
          attn_weights, p=self.attn_dropout_p, training=self.training
      )
  
      context_states = torch.matmul(attn_weights, value_states)
      context_states = (
          context_states.permute(0, 2, 1, 3)
          .contiguous()
          .view(batch_size, seq_length, -1)
      )
      attn_output = self.out_proj(context_states)
  
      return attn_output, attn_weights, past_key_value
  
imports: import math
  from typing import Optional, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
  from torch.nn import functional
  from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
  from ...modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
  from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions, QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast, TokenClassifierOutput
  from ...modeling_utils import PreTrainedModel
  from ...utils import logging
  from .configuration_mpt import MptConfig
label_desc: Use image node_id 2 for calling the MptAttention obj's underlying member method code with example usage: obj.forward(hidden_states, position_bias, past_key_value, attention_mask) and returns: attn_output,  attn_weights,  past_key_value