Add test

Author: tengyifei

test/spmd/test_xla_sharding.py

@@ -618,9 +618,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())
@@ -1383,8 +1383,9 @@ 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")
@@ -1405,8 +1406,9 @@ 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")
@@ -1427,8 +1429,9 @@ 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")
@@ -1569,6 +1572,7 @@ def test_mark_sharding_with_gradients_annotation(self):
       # Check that the gradient has sharding.
       self.assertIn(sharding_spec, x_grad_sharding)
 
+<<<<<<< HEAD
   def test_valid_mesh_creation(self):
     mesh_shape = (1, self.n_devices)
     axis_names = ('data', 'model')
@@ -1660,6 +1664,26 @@ def test_get_logical_mesh(self):
     self.assertEqual(logical_mesh.shape, mesh_shape)
     np.testing.assert_array_equal(np.sort(logical_mesh.flatten()), device_ids)
 
+  def test_shard_as(self):
+    mesh = self._get_mesh((self.n_devices,))
+    partition_spec = (0,)
+    x = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8],
+                     dtype=torch.float,
+                     device=xm.xla_device())
+    x = xs.mark_sharding_with_gradients(x, mesh, partition_spec)
+    y = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8],
+                     dtype=torch.float,
+                     device=xm.xla_device())
+
+    x, y = xs.shard_as(x, y)
+    torch_xla.sync()
+
+    sharding_spec = '{devices=[%d]' % self.n_devices
+    x_sharding = torch_xla._XLAC._get_xla_sharding_spec(x)
+    y_sharding = torch_xla._XLAC._get_xla_sharding_spec(y)
+    self.assertIn(sharding_spec, x_sharding)
+    self.assertEqual(x_sharding, y_sharding)
+
 
 if __name__ == '__main__':
   test = unittest.main()

torch_xla/distributed/spmd/__init__.py

@@ -4,7 +4,7 @@
     mark_sharding, mark_sharding_with_gradients, clear_sharding, get_1d_mesh,
     wrap_if_sharded, xla_patched_nn_linear_forward, set_global_mesh,
     get_global_mesh, _mark_manual_sharding, enable_manual_sharding,
-    disable_manual_sharding, apply_backward_optimization_barrier)
+    disable_manual_sharding, apply_backward_optimization_barrier, shard_as)
 from .api import xla_distribute_tensor, xla_distribute_module, auto_policy
 
 __all__ = [
@@ -19,6 +19,7 @@
     "MarkShardingFunction"
     "mark_sharding",
     "mark_sharding_with_gradients",
+    "shard_as",
     "clear_sharding",
     "get_1d_mesh",
     "wrap_if_sharded",