code: def get_cu_seqlens(attn_mask):
      """generate a cumulative sequence length mask for flash attention using attn mask"""
      if len(attn_mask.shape) == 1:
          attn_mask = attn_mask.unsqueeze(0)
  
      device = attn_mask.device
      results = []
      max_seq_lens = []
  
      for row in attn_mask:
          # Exclude zeros to avoid adding their positions to the mask
          t_non_zeros = row[row != 0]
          # Find where the sequence number changes (including the first position)
          seq_change = torch.cat(
              [
                  torch.tensor([1], dtype=torch.int32, device=device),
                  t_non_zeros[1:] != t_non_zeros[:-1],
              ]
          )
          # Get the indices where the sequence changes
          change_indices = torch.cat(
              [
                  (seq_change == 1).nonzero(as_tuple=True)[0],
                  torch.tensor(
                      [len(t_non_zeros)],
                      dtype=torch.int32,
                      device=device,
                  ),
              ]
          )
          # Calculate the sequence lengths
          seq_lengths = change_indices[1:] - change_indices[:-1]
          # Calculate the length of the final sequence or padding
          final_seq_length = len(row) - change_indices[-1]
          # Append the length of the final sequence or padding to seq_lengths
          if final_seq_length.item():
              seq_lengths = torch.cat(
                  [
                      seq_lengths,
                      torch.tensor(
                          [final_seq_length.item()],
                          dtype=torch.int32,
                          device=device,
                      ),
                  ]
              )
          # Calculate the cumulative sequence lengths
          cu_seqlens = torch.cat(
              [
                  torch.tensor([0], dtype=torch.int32, device=device),
                  seq_lengths.cumsum(0),
              ]
          )
          max_seq_len = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
          results.append(cu_seqlens)
          max_seq_lens.append(max_seq_len)
  
      return torch.stack(results).to(dtype=torch.int32), torch.stack(
          max_seq_lens
      )
  
imports: import torch
label_desc: Use image node_id 1 for calling a global function with example usage: get_cu_seqlens(attn_mask) and returns: torch,  torch