code: def __init__(
      self,
      dim,
      base=10000,
      distance_scale: Union[int, float] = 1,
      dtype: torch.dtype = torch.half,
  ):
      super().__init__()
      inv_freq = 1.0 / (
          base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)
      )
      inv_freq = inv_freq.to(dtype)
      self.distance_scale = distance_scale
      self.dtype = dtype
      self.inv_freq = inv_freq
  
imports: import math
  from typing import Union
  import torch
  import bmtrain
  import torch.nn.functional
label_desc: Use image node_id 1 to create a new RotaryEmbedding object from inherited base classes: bmt with example: obj = RotaryEmbedding(dim, base, distance_scale, dtype)