Get index of a 2D matrix based on values in the columns

1

Give a 2D matrix of integers, I would like to get the index of the row that has the corresponding value. For example, if I have a matrix A = [[1,2], [2,3], [3,4]] and if I have the value for column 1 =3, column 2=4, which leads to [3,4]. Then this function should return index number 2 from matrix A given that the third row corresponds to [3,4]. How can I do this in PyTorch?

2

You could try the following snippet:

a = torch.FloatTensor([[1, 2], [2, 3], [3, 4]])
b = torch.FloatTensor([[3, 4]])
(a == b).nonzero()
1 Like
3

Thank you. But this returns the following:

2  0
2  1
[torch.LongTensor of size 2x2]

Do you know if it is possible to return the row index, which is 2 in our case?

4

How about this?

number1 = 3 # look for 3 in column 1
number2 = 4
column1 = A[:,0]
column2 = A[:,1]
_, index1 = torch.min(torch.abs(column1 - number1), 0)
_, index2 = torch.min(torch.abs(column2 - number2), 0)

N.B. This isn’t differentiable.