yapf

Author: qihqi

test/test_pallas_spmd.py

@@ -84,14 +84,16 @@ def test_flash_attention_spmd_data_parallel(self):
   def test_flash_attention_spmd_data_parallel_5d(self):
     n_devices = xr.global_runtime_device_count()
     xs.set_global_mesh(
-      xs.Mesh(range(n_devices), (n_devices // 2, 2, 1, 1, 1),
-      ('fsdp', 'dp', 'a', 'b', 'c')))
+        xs.Mesh(
+            range(n_devices), (n_devices // 2, 2, 1, 1, 1),
+            ('fsdp', 'dp', 'a', 'b', 'c')))
 
     q = torch.randn(4, 2, 2, 128, 4).to("xla")
     k = torch.randn(4, 2, 2, 128, 4).to("xla")
     v = torch.randn(4, 2, 2, 128, 4).to("xla")
 
-    o = flash_attention(q, k, v, partition_spec=('fsdp', 'dp', None, None, None))
+    o = flash_attention(
+        q, k, v, partition_spec=('fsdp', 'dp', None, None, None))
     dev_ids = ','.join(map(str, range(n_devices)))
     self.assertEqual(
         torch_xla._XLAC._get_xla_sharding_spec(o),
@@ -144,7 +146,7 @@ def test_flash_attention_backward_spmd_data_parallel(self):
     q_grad = q.grad
     k_grad = k.grad
     v_grad = v.grad
-    
+
     dev_ids = ','.join(map(str, range(n_devices)))
     self.assertEqual(
         torch_xla._XLAC._get_xla_sharding_spec(q_grad),

torch_xla/experimental/custom_kernel.py

@@ -90,7 +90,9 @@ def wrapped(*args):
       for i, r in enumerate(res):
         if isinstance(r, torch.Tensor):
           assert str(r.device).startswith('xla'), f'{i}th device is {r.device}'
-          assert len(r.shape) == len(output_specs[i]), f'{i}th shape is {r.shape}, sharding is {output_specs[i]}'
+          assert len(r.shape) == len(
+              output_specs[i]
+          ), f'{i}th shape is {r.shape}, sharding is {output_specs[i]}'
       return tuple(
           xs.disable_manual_sharding(a, spec, _full_shape(a, spec), mesh=mesh).
           global_tensor
@@ -387,7 +389,8 @@ def _fa_custom_forward_single_device(
       args += [ab]
     if segment_ids is not None:
       args += [q_segment_ids_fa, kv_segment_ids_fa]
-    custom_call_output = torch_xla._XLAC._xla_tpu_custom_call(args, payload, shapes, dtypes)
+    custom_call_output = torch_xla._XLAC._xla_tpu_custom_call(
+        args, payload, shapes, dtypes)
 
     assert isinstance(custom_call_output, list)
     if not save_residuals: