code: def index_cos_sin(cos, sin, seqlen_offsets, seqlen):
      if isinstance(seqlen_offsets, torch.Tensor):
          batch_size = seqlen_offsets.shape[0]
          arange = rearrange(
              torch.arange(seqlen, device=cos.device), "s -> 1 s"
          )
          idx = rearrange(seqlen_offsets, "b -> b 1") + arange
          cos_pt = rearrange(
              cos[idx.flatten()], "(b s) d -> b s d", b=batch_size
          )
          sin_pt = rearrange(
              sin[idx.flatten()], "(b s) d -> b s d", b=batch_size
          )
      else:
          cos_pt = cos[seqlen_offsets : seqlen_offsets + seqlen]
          sin_pt = sin[seqlen_offsets : seqlen_offsets + seqlen]
      return cos_pt, sin_pt
  
file_path: flash-attention/tests/test_rotary.py
imports: import math
  import random
  import pytest
  import torch
  import torch.nn.functional
  from einops import rearrange
  from flash_attn.layers.rotary import apply_rotary_emb, apply_rotary_emb_torch
  from flash_attn.layers.rotary import apply_rotary_emb_qkv_, apply_rotary_emb_kv_
  from flash_attn.bert_padding import pad_input, unpad_input
label_desc: Use image node_id 3 for calling a global function with example usage: index_cos_sin(cos, sin, seqlen_offsets, seqlen) and returns: cos_pt,  sin_pt