tensor operations
data shuffling
machine learning
deep learning
Python programming

shuffling two tensors in the same order

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

If two tensors represent aligned data, such as features and labels, they must be shuffled with the same permutation. Shuffling them independently breaks the correspondence and silently corrupts the dataset. The correct pattern is simple: generate one index order, then apply it to both tensors.

The Core Idea

Suppose x[i] belongs with y[i]. After shuffling, those two items must still travel together. The safest method is:

  1. create a permutation of row indices
  2. reorder both tensors with that permutation

That works in NumPy, PyTorch, TensorFlow, and most other array libraries.

TensorFlow Example

Here is a direct TensorFlow solution:

python
1import tensorflow as tf
2
3x = tf.constant([[10.0], [20.0], [30.0], [40.0]])
4y = tf.constant([1, 0, 1, 0])
5
6indices = tf.random.shuffle(tf.range(tf.shape(x)[0]))
7x_shuffled = tf.gather(x, indices)
8y_shuffled = tf.gather(y, indices)
9
10print(indices.numpy())
11print(x_shuffled.numpy().tolist())
12print(y_shuffled.numpy().tolist())

The same indices tensor is used in both tf.gather calls, so row alignment is preserved.

PyTorch Example

In PyTorch, the equivalent pattern uses torch.randperm:

python
1import torch
2
3x = torch.tensor([[10.0], [20.0], [30.0], [40.0]])
4y = torch.tensor([1, 0, 1, 0])
5
6perm = torch.randperm(x.size(0))
7x_shuffled = x[perm]
8y_shuffled = y[perm]
9
10print(perm)
11print(x_shuffled)
12print(y_shuffled)

Again, one permutation drives both reorderings.

Dataset APIs Can Be Even Better

If you are already using a dataset abstraction, let it keep the pairs together for you. In TensorFlow:

python
1import tensorflow as tf
2
3x = tf.constant([[10.0], [20.0], [30.0], [40.0]])
4y = tf.constant([1, 0, 1, 0])
5
6dataset = tf.data.Dataset.from_tensor_slices((x, y))
7dataset = dataset.shuffle(buffer_size=4, seed=123)
8
9for features, label in dataset:
10    print(features.numpy(), label.numpy())

This is often cleaner in training pipelines because the pair structure is preserved from the start.

Reproducibility

If you need reproducible shuffles, fix the random seed. The exact API depends on the framework:

  • NumPy uses np.random.seed(...)
  • TensorFlow offers tf.random.set_seed(...)
  • PyTorch uses torch.manual_seed(...)

Be aware that full training reproducibility can require more than just one seed, but synchronized shuffling at least starts with a deterministic permutation.

Why Independent Shuffle Calls Fail

This is wrong:

python
1import tensorflow as tf
2
3x = tf.constant([[10.0], [20.0], [30.0], [40.0]])
4y = tf.constant([1, 0, 1, 0])
5
6x_bad = tf.random.shuffle(x)
7y_bad = tf.random.shuffle(y)

Both tensors are shuffled, but not by the same order. The data may still look random, which makes the bug easy to miss. Model accuracy then drops for mysterious reasons because the labels no longer match the inputs.

When More Than Two Tensors Are Involved

The same rule scales naturally. If you have inputs, labels, sample weights, masks, or metadata arrays, apply the same permutation to all of them:

python
1perm = tf.random.shuffle(tf.range(tf.shape(x)[0]))
2x = tf.gather(x, perm)
3y = tf.gather(y, perm)
4weights = tf.gather(weights, perm)

The permutation is the source of truth.

Common Pitfalls

The most common mistake is calling a shuffle function on each tensor separately. That breaks alignment immediately.

Another mistake is shuffling along the wrong axis. For supervised learning, you usually want to shuffle rows, not columns or feature dimensions.

A third issue is forgetting that shuffling may happen inside a data loader already. Double shuffling is not always wrong, but it can make debugging harder.

Summary

  • To shuffle paired tensors correctly, generate one permutation and apply it to both.
  • 'tf.gather and torch.randperm are the standard tools for this.'
  • Dataset APIs can keep pairs together automatically.
  • Independent shuffle calls on features and labels are incorrect.
  • Use fixed seeds when you need repeatable shuffle order.

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.