def all_gather_cpu(tree: PyTree) -> PyTree:
  """Apply multi-host all_gather operation on a tree of CPU arrays.

  Args:
    tree: Input tree containing np.ndarray in shape of (num_dev, dev_bs, ...).
  Returns:
    Output tree in the same structure as input with shape (global_bs, ...),
    where global_bs = num_host * num_dev * dev_bs.
  """
  if jax.process_count() == 1:
    return jax.tree_util.tree_map(lambda x: x.reshape(-1, *x.shape[2:]), tree)
  tree = jax.tree_util.tree_map(jnp.asarray, tree)  # (num_dev, dev_bs, ...)
  tree = jax.pmap(  # (num_dev, global_bs, ...)
      lambda x: jax.lax.all_gather(x, axis_name='device', tiled=True),
      axis_name='device',
  )(tree)
  tree = jax_utils.unreplicate(tree)  # (global_bs, ...)
  tree = jax.tree_util.tree_map(np.asarray, tree)
  return tree
Apply multi-host all_gather operation on a tree of CPU arrays.

  Args:
    tree: Input tree containing np.ndarray in shape of (num_dev, dev_bs, ...).
  Returns:
    Output tree in the same structure as input with shape (global_bs, ...),
    where global_bs = num_host * num_dev * dev_bs.
  
videogvt/train_lib/train_utils.py