is there a faster way to compare two raws in a two tensors than for loop?

user_1771445 0
Viewed 52

im new to pytorch and i would like to check for each element in the tensor1 if it is part of the same row in the tensor2 for example if i have two tensor as follwoing:

a = torch.tensor([[0,5],[1,5], [4,5], [7,8]])
b = torch.tensor([[0,1],[2,3], [7,5], [-1,7]])

the output should be as following:

[ True, False]
[False, False]
[False, False]
[False, False]
[False, False]
[False,  True]
[False,  True]
[False, False]

i know that i can do it using a simple for loop but is there a way to do it faster in pytorch?

tensor pytorch
1 Answers

you need to use the repeat_interleave for both tensors and then reshape the first tensor, and then you compare it with the other tensor like this

res = a.repeat_interleave(2, dim=1).reshape(-1, 2) == b.repeat_interleave(2, dim=0)

the output should be as follows

output
tensor([[ True, False],
        [False, False],
        [False, False],
        [False, False],
        [False, False],
        [False,  True],
        [False,  True],
        [False, False]])

the output is tensor. If you need it without [] you can use the flatten method

user_17714450
Related