class_docstr: The rotary position embeddings from RoFormer_ (Su et. al).
  A crucial insight from the method is that the query and keys are
  transformed by rotation matrices which depend on the relative positions.
  
  Other implementations are available in the Rotary Transformer repo_ and in
  GPT-NeoX_, GPT-NeoX was an inspiration
  
  .. _RoFormer: https://arxiv.org/abs/2104.09864
  .. _repo: https://github.com/ZhuiyiTechnology/roformer
  .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox
  
  If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554).
  A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96
  Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py
class_docstr_tok: ["The","rotary","position","embeddings","from","RoFormer_","(","Su","et",".","al",")",".","A","crucial","insight","from","the","method","is","that","the","query","and","keys","are","transformed","by","rotation","matrices","which","depend","on","the","relative","positions",".","Other","implementations","are","available","in","the","Rotary","Transformer","repo_","and","in","GPT-NeoX_",",","GPT-NeoX","was","an","inspiration","..","_RoFormer",":","https",":","\/\/arxiv.org\/abs\/2104.09864","..","_repo",":","https",":","\/\/github.com\/ZhuiyiTechnology\/roformer","..","_GPT-NeoX",":","https",":","\/\/github.com\/EleutherAI\/gpt-neox","If","scale_base","is","not","None",",","this","implements","XPos","(","Sun","et","al.",",","https",":","\/\/arxiv.org\/abs\/2212.10554",")",".","A","recommended","value","for","scale_base","is","512",":","https",":","\/\/github.com\/HazyResearch\/flash-attention\/issues\/96","Reference",":","https",":","\/\/github.com\/sunyt32\/torchscale\/blob\/main\/torchscale\/component\/xpos_relative_position.py"]
docstr: qkv: (batch, seqlen, 3, nheads, headdim) if kv is none,
       else it's just q of shape (batch, seqlen, nheads, headdim)
  kv: (batch, seqlen, 2, nheads, headdim)
  seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount.
      Most commonly used in inference when we have KV cache.
      If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one
      should pass in max_seqlen, which will update the cos / sin cache up to that length.
  Apply rotary embedding *inplace* to qkv and / or kv.
docstr_tok: ["qkv",":","(","batch",",","seqlen",",","3",",","nheads",",","headdim",")","if","kv","is","none",",","else","it","'s","just","q","of","shape","(","batch",",","seqlen",",","nheads",",","headdim",")","kv",":","(","batch",",","seqlen",",","2",",","nheads",",","headdim",")","seqlen_offset",":","(","batch_size",",",")","or","int",".","Each","sequence","in","x","is","shifted","by","this","amount",".","Most","commonly","used","in","inference","when","we","have","KV","cache",".","If","it","'s","a","tensor","of","shape","(","batch_size",",",")",",","then","to","update","the","cos","\/","sin","cache",",","one","should","pass","in","max_seqlen",",","which","will","update","the","cos","\/","sin","cache","up","to","that","length",".","Apply","rotary","embedding","*","inplace","*","to","qkv","and","\/","or","kv","."]
code: def forward(
      self,
      qkv: torch.Tensor,
      kv: Optional[torch.Tensor] = None,
      seqlen_offset: Union[int, torch.Tensor] = 0,
      max_seqlen: Optional[int] = None,
  ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
      """
      qkv: (batch, seqlen, 3, nheads, headdim) if kv is none,
           else it's just q of shape (batch, seqlen, nheads, headdim)
      kv: (batch, seqlen, 2, nheads, headdim)
      seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount.
          Most commonly used in inference when we have KV cache.
          If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one
          should pass in max_seqlen, which will update the cos / sin cache up to that length.
      Apply rotary embedding *inplace* to qkv and / or kv.
      """
      seqlen = qkv.shape[1]
      if max_seqlen is not None:
          self._update_cos_sin_cache(
              max_seqlen, device=qkv.device, dtype=qkv.dtype
          )
      elif isinstance(seqlen_offset, int):
          self._update_cos_sin_cache(
              seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype
          )
      if kv is None:
          if self.scale is None:
              return apply_rotary_emb_qkv_(
                  qkv,
                  self._cos_cached,
                  self._sin_cached,
                  interleaved=self.interleaved,
                  seqlen_offsets=seqlen_offset,
              )
          else:
              return apply_rotary_emb_qkv_(
                  qkv,
                  self._cos_cached,
                  self._sin_cached,
                  self._cos_k_cached,
                  self._sin_k_cached,
                  interleaved=self.interleaved,
                  seqlen_offsets=seqlen_offset,
              )
      else:
          q = qkv
          q = apply_rotary_emb_func(
              q,
              self._cos_cached,
              self._sin_cached,
              interleaved=self.interleaved,
              inplace=True,
              seqlen_offsets=seqlen_offset,
          )
          if self.scale is None:
              kv = apply_rotary_emb_kv_(
                  kv,
                  self._cos_cached,
                  self._sin_cached,
                  interleaved=self.interleaved,
                  seqlen_offsets=seqlen_offset,
              )
          else:
              kv = apply_rotary_emb_kv_(
                  kv,
                  self._cos_k_cached,
                  self._sin_k_cached,
                  interleaved=self.interleaved,
                  seqlen_offsets=seqlen_offset,
              )
          return q, kv
  
imports: import math
  from typing import Optional, Tuple, Union
  import torch
  from einops import rearrange, repeat
  from flash_attn.ops.triton.rotary import apply_rotary
label_desc: Use image node_id 4 for calling the RotaryEmbedding obj's underlying member method code with example usage: obj.forward(qkv, kv, seqlen_offset, max_seqlen) and returns: q,  kv, apply_rotary_emb_qkv_, apply_rotary_emb_qkv_