code: def opt_forward(
      self,
      hidden_states: torch.Tensor,
      key_value_states: Optional[torch.Tensor] = None,
      past_key_value: Optional[Tuple[torch.Tensor]] = None,
      attention_mask: Optional[torch.Tensor] = None,
      layer_head_mask: Optional[torch.Tensor] = None,
      output_attentions: bool = False,
      **kwargs,
  ) -> Tuple[
      torch.Tensor,
      Optional[torch.Tensor],
      Optional[Tuple[torch.Tensor]],
  ]:
      raise_on_head_mask(layer_head_mask)
  
      if output_attentions is True:
          raise ValueError(
              "output_attentions=True can not be supported with BetterTransformer."
          )
  
      # TODO: raise on batch_size = 1 + padding
  
      # if key_value_states are provided this layer is used as a cross-attention layer
      # for the decoder
      is_cross_attention = key_value_states is not None
  
      batch_size, tgt_len, _ = hidden_states.size()
  
      # get query proj
      query_states = self.q_proj(hidden_states) * self.scaling
      # get key, value proj
      if is_cross_attention and past_key_value is not None:
          # reuse k,v, cross_attentions
          key_states = past_key_value[0]
          value_states = past_key_value[1]
      elif is_cross_attention:
          # cross_attentions
          key_states = self._shape(
              self.k_proj(key_value_states), -1, batch_size
          )
          value_states = self._shape(
              self.v_proj(key_value_states), -1, batch_size
          )
      elif past_key_value is not None:
          # reuse k, v, self_attention
          key_states = self._shape(
              self.k_proj(hidden_states), -1, batch_size
          )
          value_states = self._shape(
              self.v_proj(hidden_states), -1, batch_size
          )
          key_states = torch.cat([past_key_value[0], key_states], dim=2)
          value_states = torch.cat(
              [past_key_value[1], value_states], dim=2
          )
      else:
          # self_attention
          key_states = self._shape(
              self.k_proj(hidden_states), -1, batch_size
          )
          value_states = self._shape(
              self.v_proj(hidden_states), -1, batch_size
          )
  
      if self.is_decoder:
          # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
          # Further calls to cross_attention layer can then reuse all cross-attention
          # key/value_states (first "if" case)
          # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
          # all previous decoder key/value_states. Further calls to uni-directional self-attention
          # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
          # if encoder bi-directional self-attention `past_key_value` is always `None`
          past_key_value = (key_states, value_states)
  
      query_states = self._shape(query_states, tgt_len, batch_size)
  
      query_states = query_states * self.scale
  
      dropout_p = self.dropout if self.training else 0.0
      if batch_size == 1 or self.training:
          if query_states.shape[2] > 1:
              attn_output = (
                  torch.nn.functional.scaled_dot_product_attention(
                      query_states,
                      key_states,
                      value_states,
                      attn_mask=None,
                      dropout_p=dropout_p,
                      is_causal=True,
                  )
              )
          else:
              attn_output = (
                  torch.nn.functional.scaled_dot_product_attention(
                      query_states,
                      key_states,
                      value_states,
                      attn_mask=None,
                      dropout_p=dropout_p,
                      is_causal=False,
                  )
              )
      else:
          attn_output = (
              torch.nn.functional.scaled_dot_product_attention(
                  query_states,
                  key_states,
                  value_states,
                  attn_mask=attention_mask,
                  dropout_p=dropout_p,
                  is_causal=False,
              )
          )
  
      if attn_output.size() != (
          batch_size,
          self.num_heads,
          tgt_len,
          self.head_dim,
      ):
          raise ValueError(
              f"`attn_output` should be of size {(batch_size, self.num_heads, tgt_len, self.head_dim)}, but is"
              f" {attn_output.size()}"
          )
  
      attn_output = attn_output.transpose(1, 2)
  
      # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
      # partitioned aross GPUs when using tensor-parallelism.
      attn_output = attn_output.reshape(
          batch_size, tgt_len, self.embed_dim
      )
  
      attn_output = self.out_proj(attn_output)
  
      return attn_output, None, past_key_value
  
file_path: optimum/optimum/bettertransformer/models/attention.py
imports: from typing import Optional, Tuple
  import torch
label_desc: Use image node_id 6 for calling a global function with example usage: opt_forward(hidden_states, key_value_states, past_key_value, attention_mask, layer_head_mask, output_attentions) and returns: attn_output,  None,  past_key_value