torch find indices of matching rows in 2 2D tensors
I have two 2D tensors, in different length, both are different subsets of the same original 2d tensor and I would like to find all the matching "rows"
e.g

A = [[1,2,3],[4,5,6],[7,8,9],[3,3,3]
B = [[1,2,3],[7,8,9],[4,4,4]]
torch.2dintersect(A,B) -&gt; [0,2] (the indecies of A that B also have)


I've only see numpy solutions, that use dtype as dicts, and does not work for pytorch.


Here is how I do it in numpy

arr1 = edge_index_dense.numpy().view(np.int32)
arr2 = edge_index2_dense.numpy().view(np.int32)
arr1_view = arr1.view([('', arr1.dtype)] * arr1.shape[1])
arr2_view = arr2.view([('', arr2.dtype)] * arr2.shape[1])
intersected = np.intersect1d(arr1_view, arr2_view, return_indices=True)


This answer was posted before the OP updated the question with other restrictions that changed the problem quite a bit.

TL;DR You can do something like this:

torch.where((A == B).all(dim=1))[0]




First, assuming you have:

import torch
A = torch.Tensor([[1,2,3],[4,5,6],[7,8,9]])
B = torch.Tensor([[1,2,3],[4,4,4],[7,8,9]])


We can check that A == B returns:

&gt;&gt;&gt; A == B
tensor([[ True,  True,  True],
        [ True, False, False],
        [ True,  True,  True]])


So, what we want is: the rows in which they are all True. For that, we can use the .all() operation and specify the dimension of interest, in our case 1:

&gt;&gt;&gt; (A == B).all(dim=1)
tensor([ True, False,  True])


What you actually want to know is where the Trues are. For that, we can get the first output of the torch.where() function:

&gt;&gt;&gt; torch.where((A == B).all(dim=1))[0]
tensor([0, 2])

