def rotary_embedding(
    query: torch.Tensor,
    key: torch.Tensor,
    sin: torch.Tensor,
    cos: torch.Tensor,
    rotary_dim: int,
    rotary_half: bool,
    position_ids: torch.Tensor = None,
):
    r"""
    Applies RotaryEmbedding (see https://huggingface.co/papers/2104.09864)
    on the `query ` or `key` before their multi-head attention computation.

    Args:
        query, key (torch.Tensor) : inputs to be applied with position embeddings,
            taking shape of [batch size, sequence length, num_head/num_kv_head, head_dim]
            or [num_tokens, num_head/num_kv_head, head_dim] (as well as the output shape).
        sin/cos (torch.Tensor): [num_tokens, rotary_dim] the sin/cos value tensor
            generated to be applied on query/key.
        rotary_ndims (int): the rotary dimension. e.g., 64 for GPTJ. head size for LLama.
        head_dim (int) : head dim from the input shape.
        rotary_half (bool) : if False. e.g., GPT-J 6B/ChatGLM, cos/sin is applied to the neighboring 2 elements,
            so the offset is 1.

            if True, e.g., for llama, cos/sin is applied to the neighboring rotary_dim elements,
            so the offset is rotary_dim/2.

        position_ids (torch.Tensor): Default is None and optional if sin/cos is provided.
            The according position_ids for the input. The shape should be [batch size, sequence length].

    Return
        query, key (torch.Tensor): [batch size, sequence length, num_head/num_kv_head, head_dim]
        or [num_tokens, num_head/num_kv_head, head_dim].

    """

    return RotaryEmbedding.apply_function(
        query, key, sin, cos, rotary_dim, rotary_half, position_ids
    )

    Applies RotaryEmbedding (see https://huggingface.co/papers/2104.09864)
    on the `query ` or `key` before their multi-head attention computation.

    Args:
        query, key (torch.Tensor) : inputs to be applied with position embeddings,
            taking shape of [batch size, sequence length, num_head/num_kv_head, head_dim]
            or [num_tokens, num_head/num_kv_head, head_dim] (as well as the output shape).
        sin/cos (torch.Tensor): [num_tokens, rotary_dim] the sin/cos value tensor
            generated to be applied on query/key.
        rotary_ndims (int): the rotary dimension. e.g., 64 for GPTJ. head size for LLama.
        head_dim (int) : head dim from the input shape.
        rotary_half (bool) : if False. e.g., GPT-J 6B/ChatGLM, cos/sin is applied to the neighboring 2 elements,
            so the offset is 1.

            if True, e.g., for llama, cos/sin is applied to the neighboring rotary_dim elements,
            so the offset is rotary_dim/2.

        position_ids (torch.Tensor): Default is None and optional if sin/cos is provided.
            The according position_ids for the input. The shape should be [batch size, sequence length].

    Return
        query, key (torch.Tensor): [batch size, sequence length, num_head/num_kv_head, head_dim]
        or [num_tokens, num_head/num_kv_head, head_dim].

    
intel_extension_for_pytorch/llm/functional/fusions.py