pytorch equivalent tf.gather

Question

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?

I'm not exactly sure how it's done in TF, but did you check [`torch.gather`](https://pytorch.org/docs/stable/torch.html#torch.gather)? — MBT, Dec 11 '18 at 09:35

score 6 · Accepted Answer · answered Dec 10 '18 at 19:55

6

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.

answered Dec 10 '18 at 19:55

Jatentaki

11,804
4
41
37

it results in the tensor with shape (10, 5, 30).... is it really working for OP...? – Rowing0914 Mar 04 '21 at 05:57

score 0 · Answer 2 · answered Oct 01 '21 at 06:31

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

pytorch equivalent tf.gather

2 Answers2

Exam-01

Exam-02

Linked