torch.gather和torch.take和torch.stack的等效替换

dim = 1

import torch

# torch.gather(input, dim, index, out=None),dim = 1的情况。
input = torch.tensor([[1, 2, 3], 
                      [4, 5, 6]])
index = torch.tensor([[0, 1], 
                      [1, 2]])
output_gather = torch.gather(input, 1, index)
print('torch.gather result\n',output_gather)


# 1.使用索引和广播
rows = torch.arange(input.size(0)).unsqueeze(1).expand_as(index)
'''
torch.arange(input.size(0)) = tensor([0, 1])
torch.arange(input.size(0)).unsqueeze(1) = tensor([[0],
                                                [1]])
rows = tensor([[0, 0],
               [1, 1]])
'''
output_index = input[rows, index]
print('index result\n',output_gather)


# 2.使用使用 torch.take 和 torch.flatten

linear_index = index + torch.arange(input.size(0)).unsqueeze(1) * input.size(1)
'''
tensor([[0],  * 3 
        [1]])
        
linear_index = tensor([[0, 1],
                       [4, 5]])        
'''
output = torch.take(input, linear_index)
output_3 = output.view(index.shape)
print('take and flatten result\n',output_3)
# 3.使用 torch.index_select 和 torch.stack
output_4 = torch.stack([input[i, index[i]] for i in range(input.size(0))])
print('stack result\n',output_4)
import pdb;pdb.set_trace()

dim = 0

import torch


input = torch.tensor([[1, 2, 3], 
                      [4, 5, 6]])
index = torch.tensor([[0, 1], 
                      [1, 0]])
output_gather = torch.gather(input, 0, index)
print('torch.gather result\n',output_gather)



# index1 = torch.arange(input.shape[0])
# index2 = index1.expand_as(index)

# out = input[index, index2]
# print(out)
# import pdb;pdb.set_trace()

# num_cols = input.size(1)
# linear_index = index * num_cols + torch.arange(input.size(0)).unsqueeze(0)# 加纵坐标，从0开始
# output = torch.take(input, linear_index)
# print(output)

# import pdb;pdb.set_trace()  
a= [input[index[i],i] for i in range(input.size(0))]
output_4 = torch.stack([input[index[i],i] for i in range(input.size(0))],dim=1)
print('torch.stack result\n',output_4)


import pdb;pdb.set_trace()