Function apply_rotary_inplace

pub fn apply_rotary_inplace(
    _query: &Tensor,
    _key: &Tensor,
    _cos_cache: &Tensor,
    _sin_cache: &Tensor,
    _is_neox: bool,
) -> Result<()>

Apply Rotary position encoding inplace

Arguments

• query - Query tensor of shape (num_tokens, num_heads, head_size).
• key - Key tensor of shape (num_tokens, num_kv_heads, head_size).
• cos_cache - Aligned cache of shape (num_tokens, rot_dim)
• sin_cache - Aligned cache of shape (num_tokens, rot_dim)
• is_neox - Use neox encoding instead of gpt-j style rotary