def get_local_batch(batch: Batch) -> Batch:
  """Slice local batch from an all_gathered batch."""
  global_bs = jax.tree_util.tree_leaves(batch)[0].shape[0]
  local_bs = global_bs // jax.process_count()
  proc_i = jax.process_index()
  return jax.tree_util.tree_map(
      lambda x: x[local_bs * proc_i:local_bs * (proc_i + 1)], batch)
Slice local batch from an all_gathered batch.
videogvt/train_lib/eval_utils.py