code: def rotary_pos_embs(
      q: torch.Tensor,
      k: torch.Tensor,
      seq_descs: List[DSSequenceDescriptor],
      batch: RaggedBatchWrapper,
      head_size: int,
  ):
      def make_cos_sin_emb(
          seq_len: int,
      ) -> Tuple[torch.Tensor, torch.Tensor]:
          t = torch.arange(
              seq_len,
              dtype=torch.float32,
              device=get_accelerator().current_device(),
          )
          inv_freq = (
              1.0
              / (
                  10000.0
                  ** (
                      torch.arange(
                          0,
                          head_size,
                          2,
                          dtype=torch.float32,
                          device=get_accelerator().current_device(),
                      )
                      / head_size
                  )
              )
          ).half()
  
          freqs = torch.einsum("i,j->ij", t, inv_freq)
          emb = torch.cat((freqs, freqs), dim=-1)
  
          return (
              torch.cos(emb)[:, None, :],
              torch.sin(emb)[:, None, :],
              inv_freq,
          )
  
      def rotate_half(x: torch.Tensor) -> torch.Tensor:
          return torch.cat(
              (-x[..., x.shape[-1] // 2 :], x[..., : x.shape[-1] // 2]),
              dim=-1,
          )
  
      cos, sin, freqs = make_cos_sin_emb(1024)
  
      q_out = torch.empty_like(q)
      k_out = torch.empty_like(k)
      n_heads_q = q.shape[1] // head_size
      n_heads_kv = k.shape[1] // head_size
  
      inflight_descs = batch.inflight_seq_descriptors(on_device=False)[
          : batch.current_sequences
      ]
  
      if inflight_descs.shape[0] != len(seq_descs):
          raise ValueError(
              "The number of sequence descriptors does not match the number of sequences in the batch."
          )
  
      for seq_desc, inflight_seq in zip(seq_descs, inflight_descs):
          start_idx = inflight_seq[0]
          n_tokens = seq_desc.in_flight_tokens
  
          q_src = (
              q[start_idx : start_idx + n_tokens]
              .reshape(n_tokens, n_heads_q, head_size)
              .float()
          )
          k_src = (
              k[start_idx : start_idx + n_tokens]
              .reshape(n_tokens, n_heads_kv, head_size)
              .float()
          )
          freq_start_offset = seq_desc.seen_tokens
  
          cos_chunk = cos[
              range(freq_start_offset, freq_start_offset + n_tokens)
          ]
          sin_chunk = sin[
              range(freq_start_offset, freq_start_offset + n_tokens)
          ]
  
          q_emb = q_src * cos_chunk + rotate_half(q_src) * sin_chunk
          k_emb = k_src * cos_chunk + rotate_half(k_src) * sin_chunk
  
          q_out[start_idx : start_idx + n_tokens] = q_emb.reshape(
              n_tokens, n_heads_q * head_size
          ).to(q_out.dtype)
          k_out[start_idx : start_idx + n_tokens] = k_emb.reshape(
              n_tokens, n_heads_kv * head_size
          ).to(k_out.dtype)
  
      return q_out, k_out, freqs
  
file_path: deepspeed/tests/unit/inference/v2/kernels/ragged_ops/test_blocked_rotary_emb.py
imports: from typing import List, Tuple
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.kernels.ragged_ops import BlockedRotaryEmbeddings, BlockedTrainedRotaryEmbeddings
  from deepspeed.inference.v2.ragged import RaggedBatchWrapper, DSSequenceDescriptor
  from .ragged_testing_utils import build_batch_and_manager, validate_kv_cache
  from ....v2.inference_test_utils import allclose
label_desc: Use image node_id 1 for calling a global function with example usage: rotary_pos_embs(q, k, seq_descs, batch, head_size) and returns: q_out,  k_out,  freqs, torch,  torch,  inv_freq, torch