CODEWITHSUNDEEP
tensorflowpytorch
Michael Guarino
Michael Guarino

Reputation: 61

pytorch equivalent tf.gather

I'm having some trouble porting some code over from tensorflow to pytorch.

So I have a matrix with dimensions 10x30 representing 10 examples each with 30 features. Then I have another matrix with dimensions 10x5 containing indices of the the 5 closest examples for each examples in the first matrix. I want to 'gather' using the indices contained in the second matrix the 5 closet examples for each example in the first matrix leaving me with a 3d tensor of shape 10x5x30.

In tensorflow this is done with tf.gather(matrix1, matrix2). Does anyone know how i could do this in pytorch?

Upvotes: 2

Views: 4049

Answers (2)

Shudipta Sharma
Shudipta Sharma

Reputation: 5872

I had a scenario where I had to apply gather() on an array of integers.

Exam-01

torch.Tensor().gather(dim, input_tensor)

# here,
#   input_tensor -> tensor(1)
my_list = [0, 1, 2, 3, 4]
my_tensor = torch.IntTensor(my_list)
output = my_tensor.gather(0, input_tensor) # 0 -> is the dimension

Exam-02

torch.gather(param_tensor, dim, input_tensor)

# here,
#   input_tensor -> tensor(1)
my_list = [0, 1, 2, 3, 4]
my_tensor = torch.IntTensor(my_list)
output = torch.gather(my_tensor, 0, input_tensor) # 0 -> is the dimension

Upvotes: 0

Jatentaki
Jatentaki

Reputation: 13103

How about this?

matrix1 = torch.randn(10, 30)
matrix2 = torch.randint(high=10, size=(10, 5))
gathered = matrix1[matrix2]

It uses the trick of indexing with an array of integers.

Upvotes: 7

Related Questions