Find indexes of elements from one tensor that matches in another tensor

J2M2: <p>Hi!</p>
<p>I have a unique unordered tensor A with shape (batch, N) and another unique tensor B (batch, n) where N &gt;= n.</p>
<p>I am trying to find the indexes of A that matches in B but for each batch. (batch/row wise)</p>
<p>Example</p>
<pre><code class="lang-auto">a=torch.tensor([[0,1,2,3,4,5,6,7,8,9],[6,11,1,3,14,15,9,17,18,19]])

b = torch.tensor([[2,4,6,7],[3,9,15,19]])

# weird black magic

&gt;&gt;  tensor([[2,4,6,7],[3,6,5,9]])

</code></pre>
<p>I am trying to avoid for loops because I am need to computational efficient (aka fast)</p>
<p>Thanks a lot!!!</p>

ptrblck: <p>It’s possible but you would need to:</p>
<ul>
<li>compare each value of <code>a</code> against <code>b</code> which could be faster but would need to store the intermediate results so would need more memory</li>
<li>you would then need to make sure that the number of matches for each “row” is equal in order to be able to create a single output tensor</li>
<li>you would need to undo the lexicographical sort created by <code>nonzero</code>
</li>
</ul>
<p>Here is a code snippet showing the approach.<br>
Note the comments starting with <code>!!</code> which show why you would need to undo the sort:</p>
<pre><code class="lang-python">a = torch.tensor([[0,1,2,3,4,5,6,7,8,9],[6,11,1,3,14,15,9,17,18,19]])
b = torch.tensor([[2,4,6,7],[3,9,15,19]])

idx = a.unsqueeze(2) == b.unsqueeze(1)
idx = idx.nonzero() # will sort lexicographically!
print(idx)
# tensor([[0, 2, 0],
#         [0, 4, 1],
#         [0, 6, 2],
#         [0, 7, 3],
#         [1, 3, 0],
#         [1, 5, 2], # !! due to sort !!
#         [1, 6, 1], # !! due to sort !!
#         [1, 9, 3]])

idx_ = idx[:, :2]

# check if number of matches is equal for each "row"
matches_len = idx[:,0].unique(return_counts=True)[1]
if (matches_len == matches_len[0]).all():
    output = idx[:, 1].contiguous().view(-1, matches_len[0])
    
print(output)
# tensor([[2, 4, 6, 7],
#         [3, 5, 6, 9]]) # !! 5 and 6 are sorted !!

# undo sort from nonzero
output = output[torch.arange(output.size(0)).unsqueeze(1), idx[:, 2].view_as(output)]
print(output)
# tensor([[2, 4, 6, 7],
#         [3, 6, 5, 9]])
</code></pre>

KFrank: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a></p>
<aside class="quote no-group quote-modified" data-username="ptrblck" data-post="2" data-topic="147482" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>show why you would need to undo the sort:</p>
<pre data-code-wrap="python"><code class="lang-python">idx = a.unsqueeze(2) == b.unsqueeze(1)
</code></pre>
</blockquote>
</aside>
<p>Just a quick note:</p>
<p>Unsqueezing “in the other direction,” so to speak, avoids the need to<br>
undo the sort:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; (a.unsqueeze (1) == b.unsqueeze (2)).nonzero()
tensor([[0, 0, 2],
        [0, 1, 4],
        [0, 2, 6],
        [0, 3, 7],
        [1, 0, 3],
        [1, 1, 6],
        [1, 2, 5],
        [1, 3, 9]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

ptrblck: <p>Oh, that’s neat! I was just thinking how to get rid of it and now you’ve solved it. Thanks for sharing <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

J2M2: <p>Thanks a lot, guys. It worked exactly as I expected. But it continues being like black magic for me this unsqueeze trick (: slight_smile:</p>
<p>In the end, I implemented the solution of <a class="mention" href="/u/ptrblck">@ptrblck</a> but with the suggestion of <a class="mention" href="/u/kfrank">@KFrank</a>, and works like a charm.</p>
<p>See ya,</p>