numpy
tensorflow
data-shuffling
machine-learning
tensorflow-2.0

How to shuffle two numpy datasets using TensorFlow 2.0?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

When you shuffle features and labels, they must move in exactly the same order or you destroy the training pairs. That is the central rule. In TensorFlow 2, the cleanest way to preserve alignment is to build one dataset from both arrays together and shuffle the combined dataset, or to generate one shuffled index array and apply it to both NumPy arrays.

Use tf.data.Dataset to keep pairs together

TensorFlow's dataset API is the most natural approach when the arrays are already part of a training pipeline.

python
1import numpy as np
2import tensorflow as tf
3
4features = np.array([[1, 10], [2, 20], [3, 30], [4, 40]])
5labels = np.array([0, 1, 0, 1])
6
7dataset = tf.data.Dataset.from_tensor_slices((features, labels))
8dataset = dataset.shuffle(buffer_size=len(features), reshuffle_each_iteration=True)
9dataset = dataset.batch(2)
10
11for x_batch, y_batch in dataset:
12    print(x_batch.numpy(), y_batch.numpy())

Because the arrays are packed together before shuffling, each feature row stays paired with its correct label.

Use a shared shuffled index for NumPy arrays

If you want the result back as NumPy arrays rather than as a Dataset, shuffle a shared index and gather from both arrays.

python
1import numpy as np
2import tensorflow as tf
3
4features = np.array([[1, 10], [2, 20], [3, 30], [4, 40]])
5labels = np.array([0, 1, 0, 1])
6
7indices = tf.random.shuffle(tf.range(len(features)))
8
9shuffled_features = tf.gather(features, indices).numpy()
10shuffled_labels = tf.gather(labels, indices).numpy()
11
12print(shuffled_features)
13print(shuffled_labels)

The important part is that both arrays use the same indices. That preserves correspondence while still giving you shuffled NumPy-compatible results.

Why you should not shuffle separately

This is wrong:

python
np.random.shuffle(features)
np.random.shuffle(labels)

Each shuffle call produces its own order, so the feature-label mapping is lost. Even if both arrays have the same length, the training data becomes corrupted.

Buffer size matters in dataset shuffling

When using Dataset.shuffle, the buffer_size controls how random the output can be. If the buffer is the full dataset size, the shuffle is close to a full random permutation. Smaller buffers still shuffle, but only within the window held in memory.

For small and medium in-memory arrays, buffer_size=len(features) is usually the obvious choice.

For larger streaming datasets, you may choose a smaller buffer for memory reasons, but then you should understand that the shuffle quality is more local than global.

Check lengths before shuffling

Always verify that the paired arrays have matching first-dimension lengths.

python
if len(features) != len(labels):
    raise ValueError("Features and labels must have the same number of rows")

If they differ, no shuffle strategy can preserve correct pairing because the pair structure is already invalid before you begin.

Common Pitfalls

The biggest mistake is shuffling the arrays independently. That silently breaks the relationship between inputs and targets.

Another issue is using too small a shuffle buffer and assuming you got a fully random permutation. A small buffer only gives partial shuffling behavior.

Developers also forget to validate that both arrays have the same number of samples. Misaligned lengths lead to subtle bugs or immediate errors later in the pipeline.

Finally, if you need reproducible shuffling for experiments, set the random seed on the TensorFlow side or use a deterministic shared index.

That reproducibility point matters for debugging and benchmarking. A bug that appears only under one random ordering is much easier to isolate when the shuffle can be repeated exactly.

Summary

  • Keep paired arrays together when shuffling features and labels.
  • 'tf.data.Dataset.from_tensor_slices((x, y)).shuffle(...) is the cleanest TensorFlow 2 approach.'
  • If you need arrays back, shuffle one shared index and gather from both arrays.
  • Never shuffle the two arrays independently.
  • Match the shuffle buffer size to your randomness and memory requirements.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.