def _get_batch_size(batch):
  """Returns the batch size from a batch dictionary."""
  if not batch:
    raise ValueError('Cannot get batch size for an empty batch dictionary.')

  example = next(iter(batch.values()))
  return example.shape[0]