Python PyTorch: How to Shuffle Columns or Rows of a Matrix in PyTorch
Rearranging the rows or columns of a matrix (2D tensor) is a common operation in deep learning and data preprocessing, whether you need to randomize training samples, permute feature orders, or apply data augmentation. PyTorch provides flexible indexing that makes shuffling straightforward.
This guide explains how to shuffle both columns and rows of a matrix in PyTorch, covering manual reordering, random shuffling, and combined operations.
Understanding Matrix Indexing in PyTorch
In PyTorch, a 2D tensor (matrix) is indexed using row indices (axis 0) and column indices (axis 1), both starting at 0. You can reorder rows and columns by passing tensors of indices:
tensor[row_indices][:, column_indices]
row_indices- a tensor specifying the new order of rows.column_indices- a tensor specifying the new order of columns.
By changing the order of these index tensors, you control exactly how the matrix is rearranged.
Sample Matrix
All examples use the following 3×3 tensor:
import torch
t1 = torch.tensor([
[1, 2, 3],
[5, 6, 7],
[9, 10, 11]
])
print(t1)
Output:
tensor([[ 1, 2, 3],
[ 5, 6, 7],
[ 9, 10, 11]])
Shuffling Columns
To shuffle columns, keep the row indices in their original order and provide a reordered column index tensor.
Reversing Column Order
This moves the first column to the last position and the last column to the first:
import torch
t1 = torch.tensor([
[1, 2, 3],
[5, 6, 7],
[9, 10, 11]
])
# Reverse column order: [0, 1, 2] → [2, 1, 0]
shuffled = t1[:, torch.tensor([2, 1, 0])]
print(shuffled)
Output:
tensor([[ 3, 2, 1],
[ 7, 6, 5],
[11, 10, 9]])
Each row's elements are now in reverse column order, while the rows themselves remain unchanged.
Rotating Columns Left
This shifts every column one position to the left, wrapping the first column to the end:
import torch
t1 = torch.tensor([
[1, 2, 3],
[5, 6, 7],
[9, 10, 11]
])
# Rotate columns left: col 1 → pos 0, col 2 → pos 1, col 0 → pos 2
shuffled = t1[:, torch.tensor([1, 2, 0])]
print(shuffled)
Output:
tensor([[ 2, 3, 1],
[ 6, 7, 5],
[10, 11, 9]])
When shuffling only columns, you can simplify the syntax by using t1[:, column_indices] instead of t1[torch.tensor([0, 1, 2])][:, column_indices]. Keeping the row indices as [0, 1, 2] (original order) is redundant since the slice : already selects all rows.
Shuffling Rows
To shuffle rows, provide a reordered row index tensor and keep the column indices in their original order.
Reversing Row Order
This swaps the first and last rows:
import torch
t1 = torch.tensor([
[1, 2, 3],
[5, 6, 7],
[9, 10, 11]
])
# Reverse row order: [0, 1, 2] → [2, 1, 0]
shuffled = t1[torch.tensor([2, 1, 0])]
print(shuffled)
Output:
tensor([[ 9, 10, 11],
[ 5, 6, 7],
[ 1, 2, 3]])
Rotating Rows Down
This moves each row one position down, wrapping the last row to the top:
import torch
t1 = torch.tensor([
[1, 2, 3],
[5, 6, 7],
[9, 10, 11]
])
# Rotate rows: row 1 → pos 0, row 2 → pos 1, row 0 → pos 2
shuffled = t1[torch.tensor([1, 2, 0])]
print(shuffled)
Output:
tensor([[ 5, 6, 7],
[ 9, 10, 11],
[ 1, 2, 3]])
Randomly Shuffling Rows or Columns
In practice, you often want a random shuffle rather than a specific reordering. Use torch.randperm() to generate a random permutation of indices.
Random Column Shuffle
import torch
t1 = torch.tensor([
[1, 2, 3],
[5, 6, 7],
[9, 10, 11]
])
# Generate a random permutation of column indices
col_perm = torch.randperm(t1.shape[1])
print("Column permutation:", col_perm)
shuffled = t1[:, col_perm]
print(shuffled)
Output (example - varies each run):
Column permutation: tensor([1, 2, 0])
tensor([[ 2, 3, 1],
[ 6, 7, 5],
[10, 11, 9]])
Random Row Shuffle
import torch
t1 = torch.tensor([
[1, 2, 3],
[5, 6, 7],
[9, 10, 11]
])
# Generate a random permutation of row indices
row_perm = torch.randperm(t1.shape[0])
print("Row permutation:", row_perm)
shuffled = t1[row_perm]
print(shuffled)
Output (example - varies each run):
Row permutation: tensor([2, 1, 0])
tensor([[ 9, 10, 11],
[ 5, 6, 7],
[ 1, 2, 3]])
For reproducible results, set the random seed before generating the permutation:
torch.manual_seed(42)
row_perm = torch.randperm(t1.shape[0])
This ensures the same shuffle order every time the code runs.
Shuffling Both Rows and Columns Simultaneously
You can combine row and column permutations in a single operation:
import torch
t1 = torch.tensor([
[1, 2, 3],
[5, 6, 7],
[9, 10, 11]
])
# Shuffle both rows and columns independently
row_perm = torch.randperm(t1.shape[0])
col_perm = torch.randperm(t1.shape[1])
shuffled = t1[row_perm][:, col_perm]
print("Row permutation:", row_perm)
print("Column permutation:", col_perm)
print(shuffled)
Output (example):
Row permutation: tensor([2, 1, 0])
Column permutation: tensor([2, 0, 1])
tensor([[11, 9, 10],
[ 7, 5, 6],
[ 3, 1, 2]])
Common Mistake: Using Python Lists Instead of Tensors for Indexing
A frequent error is passing a plain Python list directly for advanced indexing, which can sometimes lead to unexpected behavior or errors with certain PyTorch operations:
import torch
t1 = torch.tensor([
[1, 2, 3],
[4, 5, 6]
])
# This works but can cause issues in more complex indexing scenarios
result = t1[[1, 0]]
The recommended approach is to always use torch.tensor() for index arrays to ensure consistent behavior:
# CORRECT: explicit tensor indexing
result = t1[torch.tensor([1, 0])]
print(result)
Output:
tensor([[4, 5, 6],
[1, 2, 3]])
Quick Reference Summary
| Operation | Syntax |
|---|---|
| Shuffle columns (manual) | t1[:, torch.tensor([new_col_order])] |
| Shuffle rows (manual) | t1[torch.tensor([new_row_order])] |
| Random column shuffle | t1[:, torch.randperm(t1.shape[1])] |
| Random row shuffle | t1[torch.randperm(t1.shape[0])] |
| Shuffle both | t1[torch.randperm(t1.shape[0])][:, torch.randperm(t1.shape[1])] |
| Reproducible shuffle | Set torch.manual_seed(seed) before randperm |
Whether you need a specific reordering or a fully random shuffle, PyTorch's tensor indexing combined with torch.randperm() gives you precise and efficient control over matrix row and column arrangement.