What is the most efficient way to shuffle each row of a tensor with different shuffling order for each of the row?

1

Suppose I have a tensor of size (3,5). I need to shuffle each of the three 5 elements row independently. All the solutions that I found shuffle all the rows with the same shuffling order (eg. using torch.randperm). Is there a way to shuffle each row with independant random shuffling order, without using a for loop to shuffle each of the row one by one?

2

Try

import torch 

x = torch.rand(3, 5)
indices = torch.argsort(torch.rand(*x.shape), dim=-1)
result = x[torch.arange(x.shape[0]).unsqueeze(-1), indices]
2 Likes
3

Instead of indexing by hand, one can use the torch.gather function.

import torch 

x = torch.rand(3, 5)
indices = torch.argsort(torch.rand_like(x), dim=-1)
result = torch.gather(x, dim=-1, index=indices)
3 Likes
4

A slightly more complicated request, in case anyone is interested. I have a multi-D tensor and I am looking to shuffle the 1st dimension with different orders for the 0th dimension. Here is a solution based on the above replies.

import torch 

x = torch.rand(3, 5, 4)
indices = torch.argsort(torch.rand(*x.shape[:2]), dim=1)
result = torch.gather(x, dim=1, index=indices.unsqueeze(-1).repeat(1, 1, x.shape[-1]))