code: def unpad_input_for_concatenated_sequences(
      hidden_states, attention_mask_in_length
  ):
      """
      Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model).
      The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286).
  
      For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is:
          ```
          [
            [2, 3, 0, 0, 0, 0],
            [3, 2, 0, 0, 0, 0],
            [6, 0, 0, 0, 0, 0]
          ]
          ```
      , which refers to the 3D-attention mask:
          ```
          [
            [
              [1, 0, 0, 0, 0, 0],
              [1, 1, 0, 0, 0, 0],
              [0, 0, 1, 0, 0, 0],
              [0, 0, 1, 1, 0, 0],
              [0, 0, 1, 1, 1, 0],
              [0, 0, 0, 0, 0, 1]
            ],
            [
              [1, 0, 0, 0, 0, 0],
              [1, 1, 0, 0, 0, 0],
              [1, 1, 1, 0, 0, 0],
              [0, 0, 0, 1, 0, 0],
              [0, 0, 0, 1, 1, 0],
              [0, 0, 0, 0, 0, 1]
            ],
            [
              [1, 0, 0, 0, 0, 0],
              [1, 1, 0, 0, 0, 0],
              [1, 1, 1, 0, 0, 0],
              [1, 1, 1, 1, 0, 0],
              [1, 1, 1, 1, 1, 0],
              [1, 1, 1, 1, 1, 1]
            ]
          ]
          ```.
  
      Arguments:
          hidden_states: (batch, seqlen, ...)
          attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none.
      Return:
          hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
          cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.
          max_seqlen_in_batch: int
      """
      length = attention_mask_in_length.sum(dim=-1)
      seqlen = attention_mask_in_length.size(-1)
      attention_mask_2d = torch.arange(
          seqlen, device=length.device, dtype=length.dtype
      ).expand(len(length), seqlen) < length.unsqueeze(1)
      real_indices_idx = torch.nonzero(
          attention_mask_in_length.flatten(), as_tuple=False
      ).flatten()
      seqlens_in_batch = attention_mask_in_length.flatten()[
          real_indices_idx
      ]
      indices = torch.nonzero(
          attention_mask_2d.flatten(), as_tuple=False
      ).flatten()
      max_seqlen_in_batch = seqlens_in_batch.max().item()
      cu_seqlens = F.pad(
          torch.cumsum(
              seqlens_in_batch, dim=0, dtype=torch.torch.int32
          ),
          (1, 0),
      )
      # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
      # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
      # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
      # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
      # so we write custom forward and backward to make it a bit faster.
      return (
          index_first_axis(
              rearrange(hidden_states, "b s ... -> (b s) ..."), indices
          ),
          indices,
          cu_seqlens,
          max_seqlen_in_batch,
      )
  
file_path: flash-attention/flash_attn/bert_padding.py
imports: import torch
  import torch.nn.functional
  from einops import rearrange, repeat
label_desc: Use image node_id 2 for calling a global function with example usage: unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length) and returns: index_first_axis,  indices,  cu_seqlens,  max_seqlen_in_batch