code: def init_to_get_rotary(config, base=10000, use_elem=False):
      """
      This function initializes the rotary positional embedding, it is compatible for all models and is called in ShardFormer
      Args:
          base : calculation arg
          use_elem : activated when using chatglm-based models
      """
      config.head_dim_ = (
          config.hidden_size // config.num_attention_heads
      )
      if not hasattr(config, "rope_scaling"):
          rope_scaling_factor = 1.0
      else:
          rope_scaling_factor = (
              config.rope_scaling.factor
              if config.rope_scaling is not None
              else 1.0
          )
  
      if hasattr(config, "max_sequence_length"):
          max_seq_len = config.max_sequence_length
      elif hasattr(config, "max_position_embeddings"):
          max_seq_len = (
              config.max_position_embeddings * rope_scaling_factor
          )
      else:
          max_seq_len = 2048 * rope_scaling_factor
      base = float(base)
  
      # NTK  ref: https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
      try:
          ntk_alpha = float(os.environ.get("INFER_NTK_ALPHA", 1))
          assert ntk_alpha >= 1
          if ntk_alpha > 1:
              print(f"Note: NTK enabled, alpha set to {ntk_alpha}")
          max_seq_len *= ntk_alpha
          base = base * (
              ntk_alpha ** (config.head_dim_ / (config.head_dim_ - 2))
          )  # Base change formula
      except:
          pass
  
      n_elem = config.head_dim_
      if use_elem:
          n_elem //= 2
  
      inv_freq = 1.0 / (
          base
          ** (
              torch.arange(
                  0, n_elem, 2, device="cpu", dtype=torch.float32
              )
              / n_elem
          )
      )
      t = (
          torch.arange(
              max_seq_len + 1024 * 64, device="cpu", dtype=torch.float32
          )
          / rope_scaling_factor
      )
      freqs = torch.outer(t, inv_freq)
  
      _cos_cached = torch.cos(freqs).to(torch.float)
      _sin_cached = torch.sin(freqs).to(torch.float)
      return _cos_cached, _sin_cached
  
imports: import math
  import os
  import types
  from collections import defaultdict
  from functools import partial
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from torch_int.nn.bmm import BMM_S8T_S8N_F32T, BMM_S8T_S8N_S8T
  from transformers import PreTrainedModel
  from transformers.modeling_outputs import BaseModelOutputWithPast
  from transformers.models.llama.configuration_llama import LlamaConfig
  from transformers.models.llama.modeling_llama import LLAMA_INPUTS_DOCSTRING, LlamaAttention, LlamaDecoderLayer, LlamaMLP, LlamaRotaryEmbedding, repeat_kv, rotate_half
  from transformers.utils import add_start_docstrings_to_model_forward
  from colossalai.inference.tensor_parallel.batch_infer_state import BatchInferState
  from colossalai.kernel.triton import copy_kv_cache_to_dest, int8_rotary_embedding_fwd, smooth_llama_context_attn_fwd, smooth_token_attention_fwd
  from .base_model import BaseSmoothForCausalLM
  from .linear import W8A8B8O8Linear, W8A8BFP32O32LinearSiLU, W8A8BFP32OFP32Linear
label_desc: Use image node_id 2 for calling a global function with example usage: init_to_get_rotary(config, base, use_elem) and returns: _cos_cached,  _sin_cached