code: def get_shard_id2num_examples(num_shards, total_num_examples):
    
    # 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
docs: Return the mapping shard_id=>num_examples, assuming round-robin.