def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        """
        YuanRotaryEmbedding is equivalent to LlamaRotaryEmbedding in transformers v4.36
        """

        super().__init__()

        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (
            self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Build here to make `torch.jit.trace` work.
        self._set_cos_sin_cache(
            seq_len=max_position_embeddings,
            device=self.inv_freq.device,
            dtype=torch.get_default_dtype(),
        )

        YuanRotaryEmbedding is equivalent to LlamaRotaryEmbedding in transformers v4.36
        
tests/cpu/hf_configs/yuan/yuan_hf_model.py