1
votes

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.

1

1 Answers

2
votes

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])