# Randomly sampling equal points ensuring equal number per class

**URL:** <https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932>\
**Category:** General Discussion\
**Tags:** tfdata\
**Created:** [May 5, 2021, 8:07am UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932 "2021-05-05T08:07:33Z")\
**Posts on this page:** 19\
**Page:** 1

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [May 5, 2021, 8:07am UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/1 "2021-05-05T08:07:33Z")

</div>

Hi folks.

Currently, I have a requirement for a batch of data that should have an equal number of samples from each of the given classes.

I am implementing it using the naive way for CIFAR10:

```python
def support_sampler():
    idx_dict = dict()
    for class_id in np.arange(0, 10):
        subset_labels = sampled_labels[sampled_labels == class_id] 
        random_sampled = np.random.choice(len(subset_labels), 16)
        idx_dict[class_id] = random_sampled
    return np.concatenate(list(idx_dict.values()))

def get_support_ds():
    random_balanced_idx = support_sampler()
    temp_train, temp_labels = sampled_train[random_balanced_idx],\
        sampled_labels[random_balanced_idx]
    support_ds = tf.data.Dataset.from_tensor_slices((temp_train, temp_labels))
    support_ds = (
        support_ds
        .shuffle(BATCH_SIZE * 1000)
        .map(agumentation, num_parallel_calls=AUTO)
        .batch(BATCH_SIZE)
    )
    return support_ds

```

Is there a better way? Particularly using pure TF ops with `tf.data`?

---

<div class="post-metadata">

**Author:** ![Mark\_Daoust](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/mark_daoust/32/411_2.png) [@Mark\_Daoust](https://discuss.ai.google.dev/u/Mark_Daoust)\
**Post date:** [May 6, 2021, 1:56pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/2 "2021-05-06T13:56:29Z")

</div>

Here the approach I used was to make a dataset for each class, and then merge them.

> **[การจำแนกประเภทข้อมูลที่ไม่สมดุล  |  TensorFlow Core](https://www.tensorflow.org/tutorials/structured_data/imbalanced_data?hl=th)**

I used `sample_from_datasets` so it’s approximately equal. But you could also zip the datasets then and .map a function to stack all the zipped tensors.

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [May 6, 2021, 2:08pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/3 "2021-05-06T14:08:45Z")

</div>

Thanks Mark. I later revisted that tutorial and found out about that neat method. Solved my purpose.

I think having a separate sampler utility for `tf.data` pipelines might be better from usability standpoint.

---

<div class="post-metadata">

**Author:** ![Mark\_Daoust](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/mark_daoust/32/411_2.png) [@Mark\_Daoust](https://discuss.ai.google.dev/u/Mark_Daoust)\
**Post date:** [May 6, 2021, 2:24pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/4 "2021-05-06T14:24:18Z")

</div>

There is this “rejection resample” function:

> **[tf.data: Build TensorFlow input pipelines  |  TensorFlow Core](https://www.tensorflow.org/guide/data#rejection_resampling)**

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [May 6, 2021, 2:41pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/5 "2021-05-06T14:41:21Z")

</div>

Oh my. This is really neat. Thanks for sharing.

I need to extend the example for my use case.

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [May 7, 2021, 11:40am UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/6 "2021-05-07T11:40:05Z")

</div>

@markdaoust here’s what I tried:

```python
def class_func(image, label):
    return label

SUPPORT_BATCH_SIZE = 640

(x_train, y_train), (_, _) = tf.keras.datasets.cifar10.load_data()

sampled_idx = np.random.choice(len(x_train), 4000)
sampled_train, sampled_labels = x_train[sampled_idx], y_train[sampled_idx].squeeze()
sampled_labels = sampled_labels.astype("int32")
support_ds = tf.data.Dataset.from_tensor_slices((sampled_train, sampled_labels))

distribution = Counter(sampled_labels)
counts = np.array(list(distribution.values()))
fractions = counts/counts.sum().astype("float64")

target_distribution = np.array([0.1] * 10).astype("float64")
resampler = tf.data.experimental.rejection_resample(
    class_func, target_dist=target_distribution, initial_dist=fractions)
support_ds = support_ds.apply(resampler).batch(SUPPORT_BATCH_SIZE)

```

Here’s the root error:

```auto
TypeError: Input 'y' of 'Less' Op has type float64 that does not match type float32 of argument 'x'.

```

Any idea what I might have missed out on? Here’s [the Colab](https://colab.research.google.com/gist/sayakpaul/8c69acc936e41976eeff1e67e8c7a7e6/scratchpad.ipynb) if you wanna give it a shot.

---

<div class="post-metadata">

**Author:** ![Mark\_Daoust](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/mark_daoust/32/411_2.png) [@Mark\_Daoust](https://discuss.ai.google.dev/u/Mark_Daoust)\
**Post date:** [May 7, 2021, 12:32pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/7 "2021-05-07T12:32:46Z")

</div>

To get your code to work, replace:

```auto
fractions = counts/counts.sum().astype("float64")

target_distribution = np.array([0.1] * 10).astype("float64")

```

with:

```auto
fractions = counts/counts.sum()
fractions = fractions.astype("float32")

target_distribution = np.array([0.1] * 10).astype("float32")

```

The implementation is just being a bit careless with the dtypes.

[Here](https://github.com/tensorflow/tensorflow/blob/v2.4.1/tensorflow/python/data/experimental/ops/resampling.py#L173) it’s does a `random_ops.random_uniform([], seed=seed) < p))`.

That uniform random returns a float32. So `p` needs to be float32, or it should say `random_ops.random_uniform([], seed=seed, dtype=p.dtype)`

Or it should assert that all those arguments are float32, or cast them to float32.

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [May 7, 2021, 12:53pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/8 "2021-05-07T12:53:56Z")

</div>

That worked. Thank you.

Indeed, `dtype` part was confusing to understand.

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [May 7, 2021, 1:06pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/9 "2021-05-07T13:06:21Z")

</div>

Although the code is working fine, the distribution is not what I would expect (the expectation here is to have a uniform distribution across the labels). Here’s a batch-wise summary:

```auto
Counter({6: 73, 1: 72, 7: 71, 5: 67, 0: 65, 8: 64, 9: 63, 4: 57, 3: 55, 2: 53})
Counter({9: 74, 0: 70, 4: 70, 2: 69, 3: 68, 1: 66, 7: 62, 6: 56, 5: 53, 8: 52})
Counter({0: 75, 3: 71, 6: 70, 1: 69, 8: 64, 9: 63, 4: 63, 2: 60, 7: 55, 5: 50})
Counter({4: 74, 0: 72, 7: 72, 1: 67, 5: 66, 6: 65, 3: 63, 9: 59, 2: 52, 8: 50})
Counter({2: 78, 7: 78, 6: 75, 1: 68, 4: 62, 5: 62, 9: 56, 0: 56, 3: 55, 8: 50})

```

For 640 samples with each batch, I would expect it to give 64 per class.

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [May 8, 2021, 2:34am UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/10 "2021-05-08T02:34:25Z")

</div>

I tried another approach:

```python
sampled_idx = np.random.choice(len(x_train), 4000)
sampled_train, sampled_labels = x_train[sampled_idx], y_train[sampled_idx].squeeze()
sampled_labels = sampled_labels.astype("int32")
support_ds = tf.data.Dataset.from_tensor_slices((sampled_train, sampled_labels))

ds = []

for i in np.arange(0, 10):
    ds_label = (
        support_ds
        .filter(lambda image, label: label==i)
        .repeat())
    ds.append(ds_label)

balanced_ds = tf.data.experimental.sample_from_datasets(
    ds, [0.1] * 10).batch(SUPPORT_BATCH_SIZE)

```

But here also when I do:

```python
for samples, labels in balanced_ds.take(10):
    print(Counter(labels.numpy()))

```

the distribution does not come out as expected:

```auto
Counter({9: 74, 0: 73, 3: 71, 8: 70, 1: 70, 5: 67, 7: 64, 2: 55, 6: 51, 4: 45})
Counter({2: 76, 3: 70, 4: 68, 1: 67, 6: 64, 0: 62, 7: 62, 8: 60, 9: 56, 5: 55})
Counter({1: 78, 2: 75, 7: 74, 0: 68, 9: 67, 3: 61, 5: 58, 8: 55, 4: 54, 6: 50})
Counter({6: 82, 9: 69, 5: 68, 4: 64, 1: 63, 3: 62, 7: 62, 8: 61, 2: 56, 0: 53})
Counter({6: 76, 2: 69, 5: 69, 8: 68, 4: 67, 0: 66, 1: 59, 3: 59, 9: 55, 7: 52})
Counter({8: 77, 9: 71, 4: 68, 0: 66, 2: 66, 6: 66, 7: 64, 5: 62, 1: 60, 3: 40})
Counter({8: 86, 9: 66, 4: 65, 1: 64, 2: 62, 5: 61, 0: 60, 6: 60, 3: 58, 7: 58})
Counter({7: 75, 8: 73, 6: 70, 5: 70, 3: 68, 9: 64, 4: 61, 0: 55, 2: 53, 1: 51})
Counter({6: 78, 1: 70, 5: 67, 0: 66, 2: 66, 4: 64, 8: 60, 3: 58, 9: 56, 7: 55})
Counter({9: 75, 7: 70, 8: 69, 3: 67, 4: 65, 5: 63, 2: 62, 1: 57, 0: 57, 6: 55})

```

@markdaoust

---

<div class="post-metadata">

**Author:** ![Mark\_Daoust](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/mark_daoust/32/411_2.png) [@Mark\_Daoust](https://discuss.ai.google.dev/u/Mark_Daoust)\
**Post date:** [May 8, 2021, 1:17pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/11 "2021-05-08T13:17:13Z")

</div>

Don’t trust a person’s ability to evaluate a probability distribution at a glance.

Here’s an independent implementation that gets equivalent results:

```auto
import numpy as np

for _ in range(10):
  d = np.zeros(10)
  for n in range(640):
    d[np.random.randint(10)] += 1
  print(sorted(d, reverse=True))

```

```auto
[79.0, 71.0, 69.0, 68.0, 65.0, 61.0, 60.0, 59.0, 58.0, 50.0]
[78.0, 70.0, 70.0, 68.0, 67.0, 64.0, 62.0, 57.0, 56.0, 48.0]
[78.0, 73.0, 70.0, 69.0, 67.0, 62.0, 59.0, 57.0, 53.0, 52.0]
[74.0, 71.0, 70.0, 68.0, 66.0, 61.0, 61.0, 60.0, 56.0, 53.0]
[77.0, 70.0, 67.0, 65.0, 65.0, 63.0, 62.0, 60.0, 57.0, 54.0]
[76.0, 73.0, 68.0, 67.0, 66.0, 61.0, 59.0, 58.0, 56.0, 56.0]
[74.0, 74.0, 70.0, 69.0, 68.0, 67.0, 65.0, 59.0, 48.0, 46.0]
[85.0, 69.0, 68.0, 66.0, 62.0, 61.0, 61.0, 59.0, 56.0, 53.0]
[73.0, 71.0, 67.0, 67.0, 65.0, 63.0, 61.0, 58.0, 58.0, 57.0]
[72.0, 70.0, 68.0, 67.0, 65.0, 63.0, 62.0, 60.0, 59.0, 54.0]

```

I’m not sure what the right statistical test is (something [Dirichlet](https://en.wikipedia.org/wiki/Dirichlet_distribution).) but use a bigger sample size and you’ll see that it’s converging. with 1e6 samples everything’s within 1%:

```auto
d = np.random.randint(10, size=int(1e6))
counts, _ = np.histogram(d, bins=range(11))
counts

```

```auto
array([100254, 99351, 100098, 100162, 99747, 100369, 99793, 100247,
       100039, 99940])

```

If you you want to force exact balance then with one dataset per class you can:

```auto
import tensorflow as tf

datasets = tuple(tf.data.Dataset.from_tensors(n).repeat() for n in range(10))
zipped = tf.data.Dataset.zip(datasets)
stacked = zipped.map(lambda *args: tf.stack(args, axis=0))

stacked.element_spec

```

```auto
TensorSpec(shape=(10,), dtype=tf.int32, name=None)

```

```auto
tf.data.experimental.get_single_element(stacked.take(1))

```

```auto
<tf.Tensor: shape=(10,), dtype=int32, 
  numpy=array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)>

```

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [May 8, 2021, 1:45pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/12 "2021-05-08T13:45:21Z")

</div>

Thanks for the pointers.

---

<div class="post-metadata">

**Author:** ![Mark\_Daoust](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/mark_daoust/32/411_2.png) [@Mark\_Daoust](https://discuss.ai.google.dev/u/Mark_Daoust)\
**Post date:** [May 11, 2021, 12:32am UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/13 "2021-05-11T00:32:19Z")

</div>

Also:

> <https://github.com/tensorflow/tensorflow/commit/130e833f0eac6480aaa1ccecb70c8727b4033c83>
>
> Cast \`initial\_dist\` and \`target\_dist\` to \`float32\`.
> 
> Normally TensorFlow doesn't… cast things for you, but I think in this case it's reasonable because:
> 
> \- The cast is once per dataset, not in the inner loop.
> \- A float32 can handle numbers as small as \`1e-38\`. If your sample probability is near that small it's already effectively zero.
> 
> PiperOrigin-RevId: 373045036
> Change-Id: I19de99280d44d02010c34514ea53a7ef4e2fedc5

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [May 11, 2021, 8:44am UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/14 "2021-05-11T08:44:51Z")

</div>

@markdaoust it just keeps getting interesting:

> **[Google Colab](https://colab.research.google.com/gist/sayakpaul/edfcc7989f7e77ecfd21377f4e2ab416/scratchpad.ipynb)**

What I exactly wanted:

```auto
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})

```

Crux of the code:

```python
def dataset_for_class(i):
    i = tf.cast(i, tf.uint8)
    return support_ds.filter(lambda image, label: label == i).repeat()

sampled_idx = np.random.choice(len(x_train), 4000)
sampled_train, sampled_labels = x_train[sampled_idx], y_train[sampled_idx].squeeze()

support_ds = tf.data.Dataset.from_tensor_slices((sampled_train, sampled_labels))
stratified_ds = tf.data.Dataset.range(10).interleave(dataset_for_class, cycle_length=10) 
stratified_ds = stratified_ds.batch(640)

```

Notes:

- Dataset is CIFAR10.
- I made sure that the images getting batched are different as you would ntoice in the notebook provided above.

---

<div class="post-metadata">

**Author:** ![Mark\_Daoust](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/mark_daoust/32/411_2.png) [@Mark\_Daoust](https://discuss.ai.google.dev/u/Mark_Daoust)\
**Post date:** [May 11, 2021, 3:55pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/15 "2021-05-11T15:55:13Z")

</div>

Yeah, that interleave is basically equivalent to the zip.

```auto
def dataset_for_class(i):
    i = tf.cast(i, tf.uint8)
    return support_ds.filter(lambda image, label: label == i).repeat()

```

Just remember that if you’re splitting a dataset like that, the dataset for each class loads the whole dataset, and throws out all but `1/n` of it. So if you have a larger dataset with a larger number of classes you’ll probably want to cache each of the class-datasets (but there might also be a way to fix it with querues).

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [May 11, 2021, 4:10pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/16 "2021-05-11T16:10:26Z")

</div>

True that. Let’s just continue putting together our hacks and benchmark them. Who knows, future readers may find these incredibly useful.

On a slightly related note, as you may already know this kind of stratified sampling is pretty common for few-shot classification tasks (particularly for models like [Prototypical Networks](https://arxiv.org/abs/1703.05175)). Might be a good idea to work on a tutorial concerning this topic.

---

<div class="post-metadata">

**Author:** ![innat](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/innat/32/13542_2.png) [@innat](https://discuss.ai.google.dev/u/innat)\
**Post date:** [July 26, 2022, 1:11pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/17 "2022-07-26T13:11:05Z")

</div>

@Sayak_Paul @markdaoust  
Thanks for this insightful discussion and working workaround mentioned [here](http://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/452/14). It’s really helpful.

* * *

I’m trying to get similar output from `tf.data` API, especially while working with `tf-similarity` data sampler. For TFRecord format, it also adopts similar functions (`interleave`) from tf.data, [here](https://github.com/tensorflow/similarity/blob/f9b7f9997d/tensorflow_similarity/samplers/tfrecords_samplers.py#L117-L128). But those samples additionally require each class to execute continuously. For example: in a batch, with `num of repeated sample = 4`

```auto
[0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3, ...]

```

Is there any convenient function from `tf data` API to achieve this sorting after batching on the training pairs? Following approach might be the way but after `interleave`, I’m expecting much optimize approach. Any tips?

```auto
batch_size = 240
num_classes = 10

dataset_encode = tf.data.Dataset.range(num_classes)
dataset_encode = dataset_encode.interleave(dataset_for_class, 
                                           cycle_length=10)
dataset_encode = dataset_encode.batch(batch_size) 

dataset0 = tuple(dataset_encode.filter
                 (
                     lambda x, y: tf.equal
                     (
                         y[n], n
                     )
                 ) for n in range(num_classes)
                )
...
zipped = tf.data.Dataset.zip(dataset0)
...

```

* * *

### Update

One possible solution.

```auto
dataset = tf.data.Dataset.range(1, 6)  
dataset = dataset.interleave(
    lambda x: tf.data.Dataset.from_tensors(x).repeat(5),
    cycle_length=1, 
    block_length=3,
)
list(dataset.as_numpy_iterator())

[1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 
3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5])

```

---

<div class="post-metadata">

**Author:** ![innat](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/innat/32/13542_2.png) [@innat](https://discuss.ai.google.dev/u/innat)\
**Post date:** [July 28, 2022, 8:46pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/19 "2022-07-28T20:46:57Z")

</div>

Opened a ticket here for the slow process regarding this in the `tf.data` API.

[https://github.com/tensorflow/tensorflow/issues/56934](https://github.com/tensorflow/tensorflow/issues/56934)

---

<div class="post-metadata">

**Author:** ![Bhack](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/bhack/32/6388_2.png) [@Bhack](https://discuss.ai.google.dev/u/Bhack)\
**Post date:** [July 28, 2022, 11:10pm UTC](https://discuss.ai.google.dev/t/randomly-sampling-equal-points-ensuring-equal-number-per-class/13932/20 "2022-07-28T23:10:11Z")

</div>

It is also hard to create equal splits in TF datasets:

[https://github.com/tensorflow/datasets/issues/3502](https://github.com/tensorflow/datasets/issues/3502)