yapf

Author: tengyifei

test/spmd/test_xla_sharding.py

@@ -617,9 +617,9 @@ def test_inplace_add_with_sharding(self):
 
   # avoid calling xr.addressable_device_count here otherwise it will init the test
   # in non-spmd mode.
-  @unittest.skipIf(
-      xr.device_type() == 'CPU',
-      "sharding will be the same for both tensors on single device")
+  @unittest.skipIf(xr.device_type() == 'CPU',
+                   "sharding will be the same for both tensors on single device"
+                  )
   def test_shard_hashing(self):
     xt1 = torch.ones(2, 2).to(xm.xla_device())
     xt2 = torch.ones(2, 2).to(xm.xla_device())
@@ -1379,9 +1379,8 @@ def test_get_1d_mesh(self):
     self.assertEqual(mesh_without_name.mesh_shape,
                      (xr.global_runtime_device_count(),))
 
-  @unittest.skipUnless(
-      xr.global_runtime_device_count() > 1,
-      "Multiple devices required for dataloader sharding test")
+  @unittest.skipUnless(xr.global_runtime_device_count() > 1,
+                       "Multiple devices required for dataloader sharding test")
   def test_data_loader_with_sharding(self):
     device = torch_xla.device()
     mesh = xs.get_1d_mesh("data")
@@ -1402,9 +1401,8 @@ def test_data_loader_with_sharding(self):
         f"{{devices=[{mesh.size()},1,1,1]{','.join([str(i) for i in range(mesh.size())])}}}"
     )
 
-  @unittest.skipUnless(
-      xr.global_runtime_device_count() > 1,
-      "Multiple devices required for dataloader sharding test")
+  @unittest.skipUnless(xr.global_runtime_device_count() > 1,
+                       "Multiple devices required for dataloader sharding test")
   def test_data_loader_with_non_batch_size(self):
     device = torch_xla.device()
     mesh = xs.get_1d_mesh("data")
@@ -1425,9 +1423,8 @@ def test_data_loader_with_non_batch_size(self):
         f"{{devices=[{mesh.size()},1,1,1]{','.join([str(i) for i in range(mesh.size())])}}}"
     )
 
-  @unittest.skipUnless(
-      xr.global_runtime_device_count() > 1,
-      "Multiple devices required for dataloader sharding test")
+  @unittest.skipUnless(xr.global_runtime_device_count() > 1,
+                       "Multiple devices required for dataloader sharding test")
   def test_data_loader_with_non_batch_size_and_mini_batch(self):
     device = torch_xla.device()
     mesh = xs.get_1d_mesh("data")
@@ -1568,9 +1565,9 @@ def test_mark_sharding_with_gradients_annotation(self):
       # Check that the gradient has sharding.
       self.assertIn(sharding_spec, x_grad_sharding)
 
-  @unittest.skipIf(
-      xr.device_type() == 'CPU',
-      "sharding will be the same for both tensors on single device")
+  @unittest.skipIf(xr.device_type() == 'CPU',
+                   "sharding will be the same for both tensors on single device"
+                  )
   def test_shard_as(self):
     mesh = self._get_mesh((self.n_devices,))
     partition_spec = (0,)