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 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, rotate_half
  from transformers.utils import add_start_docstrings_to_model_forward
  from colossalai.inference.kv_cache.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 3 for calling a global function with example usage: init_to_get_rotary(config, base, use_elem) and returns: _cos_cached,  _sin_cached