code: def _init_to_get_rotary(self, base=10000):
      self.config.head_dim_ = (
          self.config.hidden_size // self.config.num_attention_heads
      )
      if not hasattr(self.config, "rope_scaling"):
          rope_scaling_factor = 1.0
      else:
          rope_scaling_factor = (
              self.config.rope_scaling.factor
              if self.config.rope_scaling is not None
              else 1.0
          )
      if hasattr(self.config, "max_sequence_length"):
          max_seq_len = self.config.max_sequence_length
      elif hasattr(self.config, "max_position_embeddings"):
          max_seq_len = (
              self.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 ** (self.head_dim_ / (self.head_dim_ - 2))
          )  # Base change formula
      except:
          pass
      n_elem = self.config.head_dim_ // 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)
  
      self._cos_cached = torch.cos(freqs).to(torch.float16).cuda()
      self._sin_cached = torch.sin(freqs).to(torch.float16).cuda()
      return
  
file_path: ColossalAI/colossalai/legacy/inference/tensor_parallel/modeling/chatglm2.py
imports: import os
  from typing import Optional, Tuple
  import torch
  from torch.nn import CrossEntropyLoss
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from colossalai.inference.tensor_parallel.batch_infer_state import BatchInferState
  from colossalai.kernel.triton.token_attention_kernel import Llama2TokenAttentionForwards
  from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel, GLMBlock, GLMTransformer, SelfAttention, split_tensor_along_last_dim
  from ._utils import copy_kv_to_mem_cache
label_desc: Use image node_id 1 for calling a global function with example usage: _init_to_get_rotary(base) without return types