apply_rotary_pos_emb gives different results between PT & Flax
## Environment info
- `transformers` version: 4.18.0.dev0
- Platform: Linux-5.13.0-1019-gcp-x86_64-with-glibc2.31
- Python version: 3.9.7
- Huggingface_hub version: 0.4.0
- PyTorch version (GPU?): 1.10.2+cu102 (False)
- Tensorflow version (GPU?): 2.7.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.4.0 (cpu)
- Jax version: 0.3.1
- JaxLib version: 0.3.0
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help
- Flax + GPT-J,: @patil-suraj
(I could try to help after finalizing some other PRs if you are busy)

## Description

The method `apply_rotary_pos_emb` in `modeling_gptj.py` and `modeling_flax_gptj.py` give different results while having the same inputs. See the next section.

While working on the #16280, this difference cause the more aggressive PT/Flax equivalence test failing with `1e-5`.
(This was not detected in #15841 because `FlaxGPTJModelTest` overwrite the equivalence test, which wasn't updated at that time.)

## To reproduce

```python
from transformers.models.gptj.modeling_gptj import apply_rotary_pos_emb as apply_rotary_pos_emb_pt
from transformers.models.gptj.modeling_flax_gptj import apply_rotary_pos_emb as apply_rotary_pos_emb_flax

import numpy as np
import torch

batch_size = 1
seq_len = 3
num_heads = 2
rotary_dim = 4

k_rot = np.ones(shape=(batch_size, seq_len, num_heads, rotary_dim), dtype=np.float32)

sin = np.array([[ 0.        ,  0.        ],
       [ 0.84147096,  0.00999983],
       [ 0.9092974 ,  0.01999867],
       [ 0.14112   ,  0.0299955 ],
       [-0.7568025 ,  0.03998933],
       [-0.9589243 ,  0.04997917],
       [-0.2794155 ,  0.059964  ]], dtype=np.float32)

cos = np.array([[ 1.        ,  1.        ],
       [ 0.54030234,  0.99995   ],
       [-0.41614684,  0.9998    ],
       [-0.9899925 ,  0.99955004],
       [-0.6536436 ,  0.9992001 ],
       [ 0.2836622 ,  0.99875027],
       [ 0.96017027,  0.99820054]], dtype=np.float32)

sin = sin[:seq_len]
cos = cos[:seq_len]

pt_output = apply_rotary_pos_emb_pt(torch.FloatTensor(k_rot), sincos=(torch.FloatTensor(sin), torch.FloatTensor(cos)))
flax_output = apply_rotary_pos_emb_flax(k_rot, sincos=(sin.reshape((batch_size, seq_len, 2)), cos.reshape((batch_size, seq_len, 2))))

pt_output = pt_output.numpy()
flax_output = np.array(flax_output)

max_diff = np.amax(np.abs(pt_output - flax_output))
print(max_diff)
```

gives

```
1.8185948
```

The expected outputs (`max_diff`) should be small, like < `1e-5`.





Thanks for reporting this, looking into it.