如何获取两个不同大小的 PyTorch 张量中相等元素的索引？

Question

Let's say I have two PyTorch tensors:

t_1d = torch.Tensor([6, 5, 1, 7, 8, 4, 7, 1, 0, 4, 11, 7, 4, 7, 4, 1])
t = torch.Tensor([4, 7])

I want to get the indices of exact match intersection between the sets for the tensor t_1d with tensor t.

Desired output of t_1d and t : [5, 12] (first index of exact intersection)

Preferably on GPU for large Tensors, so no loops or Numpy casts.

Answer 1

In general, we can check where each element in t is equal to elements in t_1d .

After that, shift back the last element by as many places as it misses from the first element (in general case, here shift by -1 ) and check whether arrays are equal:

intersection = (t_1d == t[0]) & torch.roll(t_1d == t[1], shifts=-1)
torch.where(intersection)[0] # torch.tensor([5, 12])