Supports arbitrary uniform partitioning in host-global array conversions. (apple#1029)

* Allows specifying PartitionSpec to host_to_global_device_array.

* Generalizes to arbitrary uniform partitioning.

* Addresses comments and adds mixed shape test.

Author: Mark Lee

axlearn/common/gda_test.py

@@ -6,9 +6,8 @@
 Some tests are intended to be run on TPU.
 """
 
-import itertools
-
 import jax
+import pytest
 from absl import logging
 from absl.testing import absltest, parameterized
 from jax import numpy as jnp
@@ -23,23 +22,17 @@
 )
 
 
+# TODO(markblee): Consolidate with utils_test.
 class GDATest(TestCase):
-    @parameterized.parameters(
-        itertools.product(
-            ((1, 1), (8, 1), (4, 2)),  # mesh_shape
-            (1, 16),  # per_host_batch_size
-            (DataPartitionType.FULL, DataPartitionType.REPLICATED),  # data_partition
-        )
-    )
-    def test_host_array_to_gda(self, mesh_shape, per_host_batch_size, data_partition):
+    def _test_host_array_to_gda(self, mesh_shape, per_host_batch_size, data_partition):
         logging.info(
             "mesh_shape=%s per_host_batch_size=%s data_partition=%s",
             mesh_shape,
             per_host_batch_size,
             data_partition,
         )
         if not is_supported_mesh_shape(mesh_shape):
-            return
+            pytest.skip(reason=f"Unsupported {mesh_shape=}")
         devices = mesh_utils.create_device_mesh(mesh_shape)
         if data_partition == DataPartitionType.FULL:
             global_batch_size = per_host_batch_size * jax.process_count()
@@ -65,6 +58,23 @@ def test_host_array_to_gda(self, mesh_shape, per_host_batch_size, data_partition
             self.assertIsInstance(output["x"], Tensor)
             self.assertSequenceEqual(output["x"].shape, (global_batch_size, 8))
 
+    @parameterized.product(
+        mesh_shape=[(1, 1)],
+        per_host_batch_size=[1, 16],
+        data_partition=[DataPartitionType.FULL, DataPartitionType.REPLICATED],
+    )
+    def test_host_array_to_gda_single(self, **kwargs):
+        self._test_host_array_to_gda(**kwargs)
+
+    @parameterized.product(
+        mesh_shape=[(8, 1), (4, 2)],
+        per_host_batch_size=[1, 16],
+        data_partition=[DataPartitionType.FULL, DataPartitionType.REPLICATED],
+    )
+    @pytest.mark.for_8_devices
+    def test_host_array_to_gda_multiple(self, **kwargs):
+        self._test_host_array_to_gda(**kwargs)
+
 
 if __name__ == "__main__":
     absltest.main()