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]
])

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]
])

shuffled = t1[:, torch.tensor([1, 2, 0])]
print(shuffled)

Output:

tensor([[ 2,  3,  1],
        [ 6,  7,  5],
        [10, 11,  9]])

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]
])

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]
])

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]
])

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]
])

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]])

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]
])

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]
])

result = t1[[1, 0]]

The recommended approach is to always use torch.tensor() for index arrays to ensure consistent behavior:

result = t1[torch.tensor([1, 0])]
print(result)

Output:

tensor([[4, 5, 6],
        [1, 2, 3]])

Quick Reference Summary

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.