docstr: Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
docstr_tok: ["Self-attention","(","if","key_value_states","is","None",")","or","attention","over","source","sentence","(","provided","by","key_value_states",")","."]
code: def forward(
      self,
      hidden_states,
      mask=None,
      key_value_states=None,
      position_bias=None,
      past_key_value=None,
      layer_head_mask=None,
      query_length=None,
      use_cache=False,
      output_attentions=False,
  ):
      """
      Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
      """
      # Input is (batch_size, seq_length, dim)
      # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length)
      # past_key_value[0] is (batch_size, n_heads, q_len - 1, dim_per_head)
  
      batch_size, seq_length = hidden_states.shape[:2]
  
      real_seq_length = seq_length
  
      if past_key_value is not None:
          assert (
              len(past_key_value) == 2
          ), f"past_key_value should have 2 past states: keys and values. Got {len(past_key_value)} past states"
          real_seq_length += (
              past_key_value[0].shape[2]
              if query_length is None
              else query_length
          )
  
      key_length = (
          real_seq_length
          if key_value_states is None
          else key_value_states.shape[1]
      )
  
      def shape(states):
          """projection"""
          return states.view(
              batch_size, -1, self.n_heads, self.key_value_proj_dim
          ).transpose(1, 2)
  
      def unshape(states):
          """reshape"""
          return (
              states.transpose(1, 2)
              .contiguous()
              .view(batch_size, -1, self.inner_dim)
          )
  
      def project(
          hidden_states, proj_layer, key_value_states, past_key_value
      ):
          """projects hidden states correctly to key/query states"""
          if key_value_states is None:
              # self-attn
              # (batch_size, n_heads, seq_length, dim_per_head)
              hidden_states = shape(proj_layer(hidden_states))
          elif past_key_value is None:
              # cross-attn
              # (batch_size, n_heads, seq_length, dim_per_head)
              hidden_states = shape(proj_layer(key_value_states))
  
          if past_key_value is not None:
              if key_value_states is None:
                  # self-attn
                  # (batch_size, n_heads, key_length, dim_per_head)
                  hidden_states = torch.cat(
                      [past_key_value, hidden_states], dim=2
                  )
              elif past_key_value.shape[2] != key_value_states.shape[1]:
                  # checking that the `sequence_length` of the `past_key_value` is the same as
                  # the provided `key_value_states` to support prefix tuning
                  # cross-attn
                  # (batch_size, n_heads, seq_length, dim_per_head)
                  hidden_states = shape(proj_layer(key_value_states))
              else:
                  # cross-attn
                  hidden_states = past_key_value
          return hidden_states
  
      # get query states
      query_states = shape(
          self.q(hidden_states)
      )  # (batch_size, n_heads, seq_length, dim_per_head)
  
      # get key/value states
      key_states = project(
          hidden_states,
          self.k,
          key_value_states,
          past_key_value[0] if past_key_value is not None else None,
      )
      value_states = project(
          hidden_states,
          self.v,
          key_value_states,
          past_key_value[1] if past_key_value is not None else None,
      )
  
      # compute scores
      scores = torch.matmul(
          query_states, key_states.transpose(3, 2)
      )  # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9
  
      if position_bias is None:
          if not self.has_relative_attention_bias:
              position_bias = torch.zeros(
                  (1, self.n_heads, real_seq_length, key_length),
                  device=scores.device,
                  dtype=scores.dtype,
              )
          else:
              position_bias = self.compute_bias(
                  real_seq_length, key_length, device=scores.device
              )
  
          # if key and values are already calculated
          # we want only the last query position bias
          if past_key_value is not None:
              position_bias = position_bias[
                  :, :, -hidden_states.size(1) :, :
              ]
  
          if mask is not None:
              position_bias = (
                  position_bias + mask
              )  # (batch_size, n_heads, seq_length, key_length)
  
      position_bias_masked = position_bias
  
      scores += position_bias_masked
      attn_weights = nn.functional.softmax(
          scores.float(), dim=-1
      ).type_as(
          scores
      )  # (batch_size, n_heads, seq_length, key_length)
      attn_weights = nn.functional.dropout(
          attn_weights, p=self.dropout, training=self.training
      )  # (batch_size, n_heads, seq_length, key_length)
  
      # Mask heads if we want to
      if layer_head_mask is not None:
          attn_weights = attn_weights * layer_head_mask
  
      attn_output = unshape(
          torch.matmul(attn_weights, value_states)
      )  # (batch_size, seq_length, dim)
      attn_output = self.o(attn_output)
  
      present_key_value_state = (
          (key_states, value_states)
          if (self.is_decoder and use_cache)
          else None
      )
      outputs = (
          (attn_output,) + (present_key_value_state,) + (position_bias,)
      )
  
      if output_attentions:
          outputs = outputs + (attn_weights,)
      return outputs
  
imports: import copy
  import math
  import warnings
  from typing import Optional, Tuple, Union
  from loguru import logger
  import torch
  import torch.distributed
  from torch import nn
  from torch.nn import CrossEntropyLoss
  from transformers.activations import ACT2FN
  from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, Seq2SeqLMOutput
  from transformers.modeling_utils import PreTrainedModel
  from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
  from transformers.utils import is_torch_fx_proxy
  from transformers import T5Config
  from text_generation_server.utils.layers import TensorParallelColumnLinear, TensorParallelEmbedding, TensorParallelRowLinear, TensorParallelHead
label_desc: Use image node_id 4 for calling the T5Attention obj's underlying member method code with example usage: obj.forward(hidden_states, mask, key_value_states, position_bias, past_key_value, layer_head_mask, query_length, use_cache, output_attentions) and returns: outputs, states, states, hidden_states