updated.

Author: tanyuqian

redco/predictors/utils.py

@@ -14,8 +14,7 @@
 
 import numpy as np
 import jax
-from jax.sharding import Mesh, PartitionSpec
-from jax.experimental.multihost_utils import host_local_array_to_global_array
+from jax.experimental.multihost_utils import process_allgather
 
 
 def add_idxes(examples):
@@ -49,14 +48,9 @@ def pred_step(pred_rng, params, batch, pred_fn, mesh):
 
 def process_batch_preds(batch_preds_with_idxes, mesh):
     if mesh is None:
-        global_mesh = Mesh(
-            devices=np.array(jax.devices()).reshape(
-                jax.process_count(), jax.local_device_count()),
-            axis_names=('host', 'local'))
-        batch_preds_with_idxes = host_local_array_to_global_array(
-            batch_preds_with_idxes,
-            global_mesh=global_mesh,
-            pspecs=PartitionSpec('host'))
+        batch_preds_with_idxes = jax.tree.map(
+            lambda t: t.reshape((-1,) + t.shape[2:]),
+            process_allgather(batch_preds_with_idxes))
 
     batch_preds_with_idxes = jax.tree.map(np.asarray, batch_preds_with_idxes)
     preds = jax.tree.map(