code: def get_cu_seqlens_from_pos_ids(position_ids):
      """generate a cumulative sequence length mask for flash attention using pos ids"""
      if len(position_ids.shape) == 1:
          position_ids = position_ids.unsqueeze(0)
  
      device = position_ids.device
      results = []
      max_seq_lens = []
  
      for row in position_ids:
          # Count the number of consecutive zeros from the right side
          padding_length = (
              (row == 0)
              .int()
              .flip(dims=[0])
              .cumprod(dim=0)
              .sum()
              .item()
          )
  
          # Adjust the row to exclude padding
          adjusted_row = (
              row[:-padding_length] if padding_length else row.clone()
          )
  
          # Find where the position resets to 0 (indicating a new sequence)
          seq_starts = torch.cat(
              [
                  torch.tensor([True], dtype=torch.bool, device=device),
                  adjusted_row[1:] == 0,
              ]
          )
          # Get the indices where the sequence starts
          start_indices = torch.cat(
              [
                  torch.nonzero(seq_starts).unbind(dim=1)[0],
                  torch.tensor(
                      [len(adjusted_row)],
                      dtype=torch.int32,
                      device=device,
                  ),
              ]
          )
          # Calculate the sequence lengths
          seq_lengths = start_indices[1:] - start_indices[:-1]
          # Calculate the cumulative sequence lengths
          cu_seqlens = torch.cat(
              [
                  torch.tensor([0], dtype=torch.int32, device=device),
                  seq_lengths.cumsum(0),
              ]
          )
          # Append the padding length to the cumulative sequence lengths
          if padding_length:
              cu_seqlens = torch.cat(
                  [
                      cu_seqlens,
                      torch.tensor(
                          [len(row)], dtype=torch.int32, device=device
                      ),
                  ]
              )
          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 2 for calling a global function with example usage: get_cu_seqlens_from_pos_ids(position_ids) and returns: torch,  torch