Fix format

abuelnasr0 · May 6, 2024 · 479d4e8 · 479d4e8
1 parent dbf4d05
commit 479d4e8
Showing 1 changed file with 24 additions and 10 deletions.
diff --git a/keras_nlp/src/layers/modeling/rotary_embedding_test.py b/keras_nlp/src/layers/modeling/rotary_embedding_test.py
@@ -186,13 +186,22 @@ def test_rope_scaling(self):
         # )
         # query, _ = apply_rotary_pos_emb(query, query, cos, sin)
         # print(query.transpose(1, 2))
-        # fmt: off
-        expected = [[[[ 1.000000000,  1.000000000,  1.000000000,  1.000000000],
-          [ 1.000000000,  1.000000000,  1.000000000,  1.000000000]],
-         [[ 0.398157001,  0.994987488,  1.357008100,  1.004987478],
-          [ 0.398157001,  0.994987488,  1.357008100,  1.004987478]],
-         [[-0.301168621,  0.989950180,  1.381773233,  1.009949803],
-          [-0.301168621,  0.989950180,  1.381773233,  1.009949803]]]] # fmt: on
+        expected = [
+            [
+                [
+                    [1.000000000, 1.000000000, 1.000000000, 1.000000000],
+                    [1.000000000, 1.000000000, 1.000000000, 1.000000000],
+                ],
+                [
+                    [0.398157001, 0.994987488, 1.357008100, 1.004987478],
+                    [0.398157001, 0.994987488, 1.357008100, 1.004987478],
+                ],
+                [
+                    [-0.301168621, 0.989950180, 1.381773233, 1.009949803],
+                    [-0.301168621, 0.989950180, 1.381773233, 1.009949803],
+                ],
+            ]
+        ]
 
         layer = RotaryEmbedding(scaling_factor=2.0)
         self.assertAllClose(
@@ -218,9 +227,14 @@ def test_rope_scaling_with_kv_cache(self):
         # )
         # query, _ = apply_rotary_pos_emb(query, query, cos, sin)
         # query.transpose(1, 2)
-        # fmt: off
-        expected = [[[[-1.412856817,  0.975714266, -0.061930716,  1.023709655],
-        [-1.412856817,  0.975714266, -0.061930716,  1.023709655]]]] # fmt: on
+        expected = [
+            [
+                [
+                    [-1.412856817, 0.975714266, -0.061930716, 1.023709655],
+                    [-1.412856817, 0.975714266, -0.061930716, 1.023709655],
+                ]
+            ]
+        ]
 
         layer = RotaryEmbedding(scaling_factor=5.0)
         self.assertAllClose(