def tohost(x):
  """Collect batches from all devices to host and flatten batch dimensions."""

  def single_tohost(x):
    n_device, n_batch, *remaining_dims = x.shape
    return np.array(x).reshape((n_device * n_batch,) + tuple(remaining_dims))

  return jax.tree.map(single_tohost, x)
Collect batches from all devices to host and flatten batch dimensions.
language/gscan/xattn_model/predict.py