class_docstr: GPTBigCode flash attention module. This module inherits from `GPTBigCodeAttention` as the weights of the module
  stays untouched. The only required change would be on the forward pass where it needs to correctly call the public
  API of flash attention and deal with padding tokens in case the input contains any of them.
class_docstr_tok: ["GPTBigCode","flash","attention","module",".","This","module","inherits","from","`","GPTBigCodeAttention","`","as","the","weights","of","the","module","stays","untouched",".","The","only","required","change","would","be","on","the","forward","pass","where","it","needs","to","correctly","call","the","public","API","of","flash","attention","and","deal","with","padding","tokens","in","case","the","input","contains","any","of","them","."]
code: def _upad_input(
      self,
      query_layer,
      key_layer,
      value_layer,
      attention_mask,
      query_length,
  ):
      indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(
          attention_mask
      )
      (
          batch_size,
          kv_seq_len,
          num_key_value_heads,
          head_dim,
      ) = key_layer.shape
  
      key_layer = index_first_axis(
          key_layer.reshape(
              batch_size * kv_seq_len, num_key_value_heads, head_dim
          ),
          indices_k,
      )
      value_layer = index_first_axis(
          value_layer.reshape(
              batch_size * kv_seq_len, num_key_value_heads, head_dim
          ),
          indices_k,
      )
      if query_length == kv_seq_len:
          query_layer = index_first_axis(
              query_layer.reshape(
                  batch_size * kv_seq_len, self.num_heads, head_dim
              ),
              indices_k,
          )
          cu_seqlens_q = cu_seqlens_k
          max_seqlen_in_batch_q = max_seqlen_in_batch_k
          indices_q = indices_k
      elif query_length == 1:
          max_seqlen_in_batch_q = 1
          cu_seqlens_q = torch.arange(
              batch_size + 1,
              dtype=torch.int32,
              device=query_layer.device,
          )  # There is a memcpy here, that is very bad.
          indices_q = cu_seqlens_q[:-1]
          query_layer = query_layer.squeeze(1)
      else:
          # The -q_len: slice assumes left padding.
          attention_mask = attention_mask[:, -query_length:]
          (
              query_layer,
              indices_q,
              cu_seqlens_q,
              max_seqlen_in_batch_q,
          ) = unpad_input(query_layer, attention_mask)
  
      return (
          query_layer,
          key_layer,
          value_layer,
          indices_q,
          (cu_seqlens_q, cu_seqlens_k),
          (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
      )
  
imports: import math
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from ...activations import ACT2FN
  from ...modeling_attn_mask_utils import AttentionMaskConverter
  from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions, SequenceClassifierOutputWithPast, TokenClassifierOutput
  from ...modeling_utils import PreTrainedModel
  from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_10, logging
  from .configuration_gpt_bigcode import GPTBigCodeConfig
label_desc: Use image node_id 4 for calling the GPTBigCodeFlashAttention2 obj's underlying member method code with example usage: obj._upad_input(query_layer, key_layer, value_layer, attention_mask, query_length) and returns: query_layer,  key_layer,  value_layer,  indices_q,  cu_seqlens_q,  cu_seqlens_k,  max_seqlen_in_batch_q,  max_seqlen_in_batch_k