adjust the code format in test_pallas.py and test_pallas_spmd.py by y…

…apf 0.30.0.
pytorch · Feb 7, 2025 · 0a639e8 · 0a639e8
1 parent 5d583a6
commit 0a639e8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/test/test_pallas.py b/test/test_pallas.py
@@ -41,10 +41,10 @@ class PallasTest(parameterized.TestCase):
   # therefore we use != instead of ==.
   def _make_attention_mask_from_segment_ids(self, q_segment_ids,
                                             kv_segment_ids):
-    return q_segment_ids.view(q_segment_ids.shape[0], 1, q_segment_ids.shape[1],
-                              1) != kv_segment_ids.view(kv_segment_ids.shape[0],
-                                                        1, 1,
-                                                        kv_segment_ids.shape[1])
+    return q_segment_ids.view(q_segment_ids.shape[0], 1,
+                              q_segment_ids.shape[1], 1) != kv_segment_ids.view(
+                                  kv_segment_ids.shape[0], 1, 1,
+                                  kv_segment_ids.shape[1])
 
   def _attention(self, q, k, v, *, attn_mask=None, ab=None):
     attn_weight = q @ k.transpose(-2, -1)

diff --git a/test/test_pallas_spmd.py b/test/test_pallas_spmd.py
@@ -41,10 +41,10 @@ class PallasTest(unittest.TestCase):
   # therefore we use != instead of ==.
   def _make_attention_mask_from_segment_ids(self, q_segment_ids,
                                             kv_segment_ids):
-    return q_segment_ids.view(q_segment_ids.shape[0], 1, q_segment_ids.shape[1],
-                              1) != kv_segment_ids.view(kv_segment_ids.shape[0],
-                                                        1, 1,
-                                                        kv_segment_ids.shape[1])
+    return q_segment_ids.view(q_segment_ids.shape[0], 1,
+                              q_segment_ids.shape[1], 1) != kv_segment_ids.view(
+                                  kv_segment_ids.shape[0], 1, 1,
+                                  kv_segment_ids.shape[1])
 
   def _attention(self, q, k, v, *, attn_mask=None, ab=None):
     attn_weight = q @ k.transpose(-2, -1)