Add additional test cases for ConvertV2ShardingToV1Test

Author: hshahTT

test/spmd/test_spmd_debugging.py

@@ -846,6 +846,11 @@ def test_tiled_sharding(self):
         [str(i) for i in range(self.n_devices)]))
     self.run_test()
 
+    self.partition_spec = (1, 0)
+    self.expected_str = '{devices=[%d,1]%s}' % (self.n_devices, ','.join(
+        [str(i) for i in range(self.n_devices)]))
+    self.run_test()
+
   @unittest.skipIf(xr.global_runtime_device_count() < 2,
                    f"Requires at least 2 devices.")
   def test_tupled_tiled_sharding(self):
@@ -873,6 +878,11 @@ def test_partial_replication_sharding(self):
         self.n_devices // 2, ','.join(str(x) for x in range(self.n_devices)))
     self.run_test()
 
+    self.partition_spec = (None, 0)
+    self.expected_str = '{devices=[1,2,%d]%s last_tile_dim_replicate}' % (
+        self.n_devices // 2, ','.join(str(x) for x in range(self.n_devices)))
+    self.run_test()
+
   @unittest.skipIf(xr.global_runtime_device_count() < 4,
                    f"Requires at least 4 devices.")
   def test_tupled_partial_replication_sharding(self):