Making apply_rotary_pos_emb generic

[xuzifei-dmatrix]

Feature request

Is it possible to move the definition of apply_rotary_pos_emb function out of the model specific files (eg: modeling_llama.py) to a more generic file (eg: utils)? I noticed they are defined repeatedly for multiple models.

Motivation

I'm trying to apply torch.fx.wrap on the apply_rotary_pos_emb function. But since they are defined in multiple model files, I need to wrap it for multiple times (eg: I would need to wrap both transformers.models.llama.modeling_llama.apply_rotary_pos_emb and transformers.models.gemma.modeling_gemma.apply_rotary_pos_emb). And I need to be aware of what new models are using apply_rotary_pos_emb and wrap their definition of the function accordingly.

Your contribution

If you could point me to the location to place the generic apply_rotary_pos_emb. I could help to create the PR. Thanks!