Return the mapping shard_id=>num_examples, assuming round-robin.

def get_shard_id2num_examples(num_shards, total_num_examples):
  """Return the mapping shard_id=>num_examples, assuming round-robin."""
  # TODO(b/130353071): This has the strong assumption that the shards have
  # been written in a round-robin fashion. This assumption does not hold, for
  # instance, with Beam generation. The mapping shard_id=>num_examples
  # should be computed during generation.

  # Minimum number of example per shards
  num_example_in_shard = total_num_examples // num_shards
  shard_id2num_examples = [num_example_in_shard for _ in range(num_shards)]
  # If there are remaining examples, we add them to the first shards
  for shard_id in range(total_num_examples % num_shards):
    shard_id2num_examples[shard_id] += 1
  return shard_id2num_examples