Equal number of samples per class in each split

I have a custom dataset with 500 samples in 10 classes. The dataset is balanced, i.e. there are 50 samples per class. I have a helper function to verify that fact:

def generate_stats(ds):
    stats = {}
    for sample in ds:
        label = sample['label'].numpy()

        if label in stats:
          stats[label] += 1
        else:
          stats[label] = 1
    print(collections.OrderedDict(sorted(stats.items())))

generate_stats(tfds.load('mydataset', split='train'))

> OrderedDict([(0, 50), (1, 50), (2, 50), (3, 50), (4, 50), (5, 50), (6, 50), (7, 50), (8, 50), (9, 50)])

Now I want to have two equal splits, i.e. the number of samples per class is the same in each split, but it doesn't work:

d1, d2 = tfds.load('mydataset', split=['train[:50%]', 'train[50%:]'])
generate_stats(d1)
generate_stats(d2)
> OrderedDict([(0, 24), (1, 20), (2, 24), (3, 27), (4, 28), (5, 19), (6, 26), (7, 28), (8, 25), (9, 29)])
> OrderedDict([(0, 26), (1, 30), (2, 26), (3, 23), (4, 22), (5, 31), (6, 24), (7, 22), (8, 25), (9, 21)])

How can I get the same number of samples per class in each split?

Start by reading the tfds.load split expressions shown in the issue and inspect how the train split is partitioned. Determine whether equal per-class counts require a new stratified-splitting capability or a documented usage pattern; done means both resulting splits contain the same count for every label.