How could I rearrange order along certain dimension?


Suppose I have a tensor a of shape (2, 3, 4), and my order tensor idx is (2, 4). I need to rearrange it along the dim=2 dimension. Such that a[0] = a[0][:,idx[0]], a[1] = a[1][:,idx[1]].
My problem is that my tensor in reality is larger than the tensor in the example, and I would like to do it in single line which is faster and requires less memory.

How could I do this please ?

perhaps something like

a.gather(2, idx.unsqueeze(1).expand(-1,3,-1))

will work

Thanks !! What is the memory overhead of this method ? Will the memory usage gets doubled due the operation of idx.unsqueeze(1).expand(-1, 3, 1)?

expand() doesn’t allocate memory, it does stride manipulations to create self-overlaps

Thanks a lot for replying !!!