rebased. fixed two tests. make style

Author: imangohari1

tests/transformers/tests/models/falcon/test_modeling_falcon.py

@@ -454,11 +454,12 @@ def test_model_rope_scaling(self):
         torch.testing.assert_close(original_sin_short, original_sin_long[:short_input_length, :])
         # Sanity check linear RoPE scaling
         # New position "x" should match original position with index "x/scaling_factor"
-        linear_scaling_rope = FalconLinearScalingRotaryEmbedding(
+        linear_scaling_rope = FalconRotaryEmbedding(
             head_dim,
             max_position_embeddings=config.max_position_embeddings,
             base=config.rope_theta,
             scaling_factor=scaling_factor,
+            rope_type="linear",
         ).to(torch_device)
         linear_cos_short, linear_sin_short = linear_scaling_rope(x, short_input_length)
         linear_cos_long, linear_sin_long = linear_scaling_rope(x, long_input_length)
@@ -471,11 +472,12 @@ def test_model_rope_scaling(self):
         # Sanity check Dynamic NTK RoPE scaling
         # Scaling should only be observed after a long input is fed. We can observe that the frequencies increase
         # with scaling_factor (or that `inv_freq` decreases)
-        ntk_scaling_rope = FalconDynamicNTKScalingRotaryEmbedding(
+        ntk_scaling_rope = FalconRotaryEmbedding(
             head_dim,
             max_position_embeddings=config.max_position_embeddings,
             base=config.rope_theta,
             scaling_factor=scaling_factor,
+            rope_type="dynamic",
         ).to(torch_device)
         ntk_cos_short, ntk_sin_short = ntk_scaling_rope(x, short_input_length)
         ntk_cos_long, ntk_sin_long = ntk_scaling_rope(x, long_input_length)

tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py

@@ -367,11 +367,12 @@ def test_model_rope_scaling(self):
         torch.testing.assert_close(original_sin_short, original_sin_long[:short_input_length, :])
         # Sanity check linear RoPE scaling
         # New position "x" should match original position with index "x/scaling_factor"
-        linear_scaling_rope = GPTNeoXLinearScalingRotaryEmbedding(
+        linear_scaling_rope = GPTNeoXRotaryEmbedding(
             head_dim,
             max_position_embeddings=config.max_position_embeddings,
             base=config.rotary_emb_base,
             scaling_factor=scaling_factor,
+            rope_type="linear",
         ).to(torch_device)
         linear_cos_short, linear_sin_short = linear_scaling_rope(x, short_input_length)
         linear_cos_long, linear_sin_long = linear_scaling_rope(x, long_input_length)
@@ -384,11 +385,12 @@ def test_model_rope_scaling(self):
         # Sanity check Dynamic NTK RoPE scaling
         # Scaling should only be observed after a long input is fed. We can observe that the frequencies increase
         # with scaling_factor (or that `inv_freq` decreases)
-        ntk_scaling_rope = GPTNeoXDynamicNTKScalingRotaryEmbedding(
+        ntk_scaling_rope = GPTNeoXRotaryEmbedding(
             head_dim,
             max_position_embeddings=config.max_position_embeddings,
             base=config.rotary_emb_base,
             scaling_factor=scaling_factor,
+            rope_type="dynamic",
         ).to(torch_device)
         ntk_cos_short, ntk_sin_short = ntk_scaling_rope(x, short_input_length)
         ntk_cos_long, ntk_sin_long = ntk_scaling_rope(x, long_input_length)