Fast method for checking for inclusion

username: dkoutsou
cooked: <p>Hey,</p>
<p>I want to check if some elements of a tensor are included in some other tensor.</p>
<p>I have come up with the following two ways to do this. First some reordering of the tensors:</p>
<pre><code class="lang-auto">x = x.view(1, -1)
y = y.view(-1, 1)
</code></pre>
<p>Then one way would be to just do this:</p>
<pre><code class="lang-auto">return (x == y).nonzero().size(0)
</code></pre>
<p>However it’s more cache efficient if the y tensor is large to chuck it like this:</p>
<pre><code class="lang-auto">chunks = y.chunk(num_chunks)
</code></pre>
<p>And then do a for loop:</p>
<pre><code class="lang-auto">count = 0
for chunk in chunks:
     count += (x == chunk).sum()
</code></pre>
<p>Both of these methods are too slow for what I want to do. For 1000 to 100000 elements the fastest implementation of both takes about 0.14s and I want it to be under 0.01s. Is there another way to do what I want? I think the <code>==</code> operator is not so efficient as it has to materialize a large array in the meantime.</p>
<p>Thank you very much!</p>

username: googlebot
cooked: <p>resort to numpy maybe? (<a href="https://numpy.org/doc/stable/reference/generated/numpy.intersect1d.html#numpy.intersect1d" rel="nofollow noopener">intersect1d</a>)</p>

username: dkoutsou
cooked: <p>That will make me lose the GPU advantage <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: googlebot
cooked: <p>What advantage? AFAIK pytorch does GPU binary search with ops like sort, topk, unique, median, that are not helpful for your task. And you say that brute force is too slow.</p>
<p>Well, unique() can tell you overlap size: (num_unique(a)+num_unique(b)) -  num_unique(cat(a,b)). Required sortings may still be too heavy though.</p>

username: dkoutsou
cooked: <p>If I use numpy then I won’t be able to run my code on GPU. That’s what I meant. Also <code>intersect1d</code> is not exactly what I want. In case of multiple matches of an element, I want all of them to be returned and not just one of them.</p>

username: googlebot
cooked: <aside class="quote no-group" data-username="dkoutsou" data-post="5" data-topic="88539">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/d/f1d935/48.png" class="avatar"> dkoutsou:</div>
<blockquote>
<p>If I use numpy then I won’t be able to run my code on GPU</p>
</blockquote>
</aside>
<p>You’d just have to temporarily copy tensor(s) to cpu, something like:</p>
<pre><code class="lang-auto">npx = x.detach().cpu().numpy()
npy = numpy_func(npx)
y = torch.from_numpy(npy).to(device=x.device)
</code></pre>
<p>overhead of this may be acceptable, YMMV</p>
<aside class="quote no-group" data-username="dkoutsou" data-post="5" data-topic="88539">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/d/f1d935/48.png" class="avatar"> dkoutsou:</div>
<blockquote>
<p>Also <code>intersect1d</code> is not exactly what I want.</p>
</blockquote>
</aside>
<p>Check other set/sorting functions there.</p>