[FlaxGPTJ] Fix bug in rotary embeddings (#16298)

c77092a5 · Suraj Patil · GitHub · 4b277483 · c77092a5
Unverified Commit c77092a5 authored Mar 21, 2022 by Suraj Patil Committed by GitHub Mar 21, 2022
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/transformers/models/gptj/modeling_flax_gptj.py src/transformers/models/gptj/modeling_flax_gptj.py +1 -1

No files found.
--- a/src/transformers/models/gptj/modeling_flax_gptj.py
+++ b/src/transformers/models/gptj/modeling_flax_gptj.py
@@ -122,7 +122,7 @@ def create_sinusoidal_positions(num_pos, dim):


 def rotate_every_two(tensor):
-    rotate_half_tensor = jnp.stack((tensor[:, :, :, 1::2], tensor[:, :, :, ::2]), axis=-1)
+    rotate_half_tensor = jnp.stack((-tensor[:, :, :, 1::2], tensor[:, :, :, ::2]), axis=-1)
    rotate_half_tensor = rotate_half_tensor.reshape(rotate_half_tensor.shape[:-2] + (-1,))
    return rotate_half_tensor