Pytorch tensor - How to get the indexes by a specific tensor

Question

I have a tensor

t = torch.tensor([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]])

and a query tensor

q = torch.tensor([1, 0, 0, 0])

Is there a way to get the indexes of q like

indexes = t.index(q) # get back [0, 3]

in pytorch?

Answer 1

How about

In [1]: torch.nonzero((t == q).sum(dim=1) == t.size(1))
Out[1]: 
tensor([[ 0],
        [ 3]])

Comparing t == q performs element-wise comparison between t and q, since you are looking for entire row match, you need to .sum(dim=1) along the rows and see what row is a perfect match == t.size(1).

---

As of v0.4.1, torch.all() supports dim argument:

torch.all(t==q, dim=1)

Answer 2

Please try this, I do not have torch installed on this PC.

import torch

t = torch.tensor([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]])
q = torch.tensor([1, 0, 0, 0])

index = torch.nonzero(torch.sum((t == q), dim=1) == t.shape[1])

Edit note: edited for the issue raised by Shai.