Feature request

Hi Huggingface Team,
We would like to request to make the 'unsqueeze(1)' in these two lines parameterized. To be precise, we would like to request that the following lines

def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
    cos = cos[position_ids].unsqueeze(1) 
    sin = sin[position_ids].unsqueeze(1)
    ...

be converted to something like the following

def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
    ...

Motivation

We are trying to import and use the apply_rotary_pos_emb function and the LlamaRotaryEmbedding class. However, our query and key tensors have the shape [batch_size, sequence_len, heads, dim]. To make them compatible with the apply_rotary_pos_emb function, we have to transpose the tensors to [batch_size, heads, sequence_len, dim], call apply_rotary_pos_emb on them and then transpose the tensors back to [batch_size, sequence_len, heads, dim]. These unnecessary transposes could be avoided if the apply_rotary_pos_emb function had a parameter which controlled the dimension on which the unsqueeze's here were applied.

Please note that the Llama Huggingface code also does similar back and forth transposes, and hence could benefit from this very small code change as well.

Your contribution

We are willing to submit a PR for this.