MLX Data take only subset / fraction of data

[FNG-2002]

Hello everyone,

I'm using MLX-Data to load the cifar10 dataset. I would like to take only a small fraction of the data during testing/development. How would I do that?

Here is my function for this:

def get_data(batch_size, num_clients=1, partition_id=0, root=None):
    tr = load_cifar10(root=root)

    tr = tr.partition(num_partitions=num_clients, partition=partition_id)

    mean = np.array([0.485, 0.456, 0.406]).reshape((1, 1, 3))
    std = np.array([0.229, 0.224, 0.225]).reshape((1, 1, 3))

    def normalize(x):
        x = x.astype("float32") / 255.0
        return (x - mean) / std

    tr_iter = (
        tr.shuffle()
        .to_stream()
        .image_random_h_flip("image", prob=0.5)
        .pad("image", 0, 4, 4, 0.0)
        .pad("image", 1, 4, 4, 0.0)
        .image_random_crop("image", 32, 32)
        .key_transform("image", normalize)
        .batch(batch_size)
        .prefetch(4, 4)
    )

    test = load_cifar10(root=root, train=False)
    test_iter = test.to_stream().key_transform("image", normalize).batch(batch_size)

    return tr_iter, test_iter

Thanks in advance ^^