[AutoParallel] fix process_mesh (#46583)

python/paddle/distributed/auto_parallel/process_mesh.py

@@ -168,7 +168,10 @@ class ProcessMesh(object):
         else:
             new_mesh = self._mesh[index]
             new_dim_names = self._dim_names[1:]
-            return ProcessMesh(new_mesh, new_dim_names)
+            if new_mesh.shape:
+                return ProcessMesh(new_mesh, new_dim_names)
+            else:
+                return ProcessMesh([new_mesh])
 
     def __enter__(self):
         set_current_process_mesh(self)

python/paddle/fluid/tests/unittests/auto_parallel/test_process_mesh.py

@@ -101,6 +101,12 @@ class TestProcessMesh(unittest.TestCase):
         self.assertEqual(sub_process_mesh4.dim_names, ["d0"])
         self.assertEqual(sub_process_mesh4.ndim, 1)
 
+        sub_process_mesh5 = sub_process_mesh3[0]
+        self.assertEqual(sub_process_mesh5.shape, [1])
+        self.assertEqual(sub_process_mesh5.process_ids, [1])
+        self.assertEqual(sub_process_mesh5.dim_names, ["d0"])
+        self.assertEqual(sub_process_mesh5.ndim, 1)
+
     def test_context_manager(self):
         mesh = np.array([1, 2, 3, 4])
         input = static.data(name="input",