Match the behavior on single host wrt multi-host if tiled=False. Fixes #25783

PiperOrigin-RevId: 713398173

Author: yashk2810

jax/experimental/multihost_utils.py

@@ -103,7 +103,8 @@ def _handle_array_process_allgather(inp, tiled):
   else:
     # All inputs here will be fully addressable.
     if jax.process_count() == 1:
-      return np.asarray(inp)
+      out = np.asarray(inp)
+      return np.expand_dims(out, axis=0) if not tiled else out
 
     devices = np.array(jax.devices()).reshape(jax.process_count(),
                                               jax.local_device_count())

tests/array_test.py

@@ -700,10 +700,14 @@ def test_array_is_ready(self):
 
   def test_process_allgather_single_host(self):
     x = jnp.arange(8.)
-    out = multihost_utils.process_allgather(x)
+    out = multihost_utils.process_allgather(x, tiled=True)
     self.assertEqual(out.shape, x.shape)
     self.assertArraysEqual(out, x)
 
+    out = multihost_utils.process_allgather(x)
+    self.assertEqual(out.shape, (1, x.shape[0]))
+    self.assertArraysEqual(out, np.expand_dims(x, axis=0))
+
   @jtu.sample_product(
     dtype=jtu.dtypes.all,
     shape=[(), (10), (2, 3)],