Update test_mark_shard_scalar to reflect mesh_shape implementation

Author: aws-cph

test/spmd/test_xla_sharding.py

@@ -1162,9 +1162,7 @@ def test_mark_shard_scalar(self):
       self.assertIsInstance(shard.indices, type(Ellipsis))
       self.assertEqual(shard.replica_id, i)
 
-    # It looks like mesh_shape attribute is never implemented.
-    with self.assertRaises(AttributeError):
-      xt.mesh_shape
+    # Removed mesh shape assert now that it has been implemented
 
   def test_global_mesh(self):
     expected_mesh = self._get_mesh((1, self.n_devices))