def _split_batch_iter(orig_iter: Iterable[Batch], batch_size: int):
  """Split batches from data loader into smaller batch size."""
  num_dev, num_splits = None, None
  for batch in orig_iter:
    if num_dev is None:
      num_dev, orig_dev_bs = jax.tree_util.tree_leaves(batch)[0].shape[:2]
      orig_batch_size = num_dev * orig_dev_bs * jax.process_count()
      assert orig_batch_size % batch_size == 0
      num_splits = orig_batch_size // batch_size
    splitted_batch = jax.tree_util.tree_map(
        lambda x: jnp.reshape(x, (num_dev, num_splits, -1, *x.shape[2:])),
        batch)
    for i in range(num_splits):
      if 'batch_mask' in splitted_batch and splitted_batch[
          'batch_mask'][:, i].sum() == 0:
        continue
      yield jax.tree_util.tree_map(lambda x: x[:, i], splitted_batch)  # pylint: disable=cell-var-from-loop
Split batches from data loader into smaller batch size.
videogvt/train_lib/train_utils.py