How to efficiently retrieve all values in tensor B, corresponding to the same value in tensor A

damiaanr: <p>Let’s say I have two tensors:</p>
<pre><code class="lang-auto">samples = torch.tensor([50, 60, 40, 50, 50, 60])
labels = torch.tensor([0, 0, 1, 2, 1, 1])
</code></pre>
<p>What I would like to do is to efficiently retrieve the values in <strong>labels</strong> for each unique value in <strong>samples</strong>. So the expected output would be three tensors (assuming ordered samples, such as when using torch.unique):</p>
<pre><code class="lang-auto">torch.tensor([1])
torch.tensor([0, 2, 1])
torch.tensor([0, 1])
</code></pre>
<p>Now, of course, I could do torch.unique(samples) to get the unique samples, and use this in a loop to index labels, for example <strong>labels[samples == 50]</strong> yields the expected <strong>torch.tensor([0, 2, 1])</strong>. However, I am not sure if that is efficient as I am still using a for loop. Should I do something smart with the return_inverse parameter of torch.unique? I haven’t solved the puzzle yet!</p>
<p>A probably even more inefficient way would be something like this (this is what I currently do):</p>
<pre><code class="lang-auto">occurrences = defaultdict(list)

for sample, label in zip(samples, labels):
    key = sample.flatten().tolist()
    occurrences[tuple(key)].append(label)
</code></pre>
<p>In my specific problem, <strong>samples</strong> is either a Python list of 2D tensors or a large 3D tensor (based on whether the tensors could be stacked, since the last dimension may vary depending on the configuration). Why I would like to collect the occurrences in <strong>labels</strong> for every unique sample is because I would like to calculate the entropy of these labels. So, currently, after the above code, I have something like this:</p>
<pre><code class="lang-auto">for sample, label_occurrences in occurrences.items():
  label_occurrences = torch.tensor(label_occurrences)
  _, counts = torch.unique(label_occurrences, return_counts=True)
  etr = entropy(counts / label_occurrences.shape[0], base=num_classes)
</code></pre>
<p>This also feels very inefficient. <strong>How can I make my code more efficient(/vectorised)?</strong></p>
<p>(Bonus: I am using SciPy’s entropy function in the last snippet; probably there are even faster ways to calculate these numbers using Torch’s built-in operations?)</p>

KFrank: <p>Hi Damiaan!</p>
<aside class="quote no-group quote-modified" data-username="damiaanr" data-post="1" data-topic="183270" 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/damiaanr/48/61196_2.png" class="avatar"> damiaanr:</div>
<blockquote>
<p>Let’s say I have two tensors:</p>
<pre><code class="lang-auto">samples = torch.tensor([50, 60, 40, 50, 50, 60])
labels = torch.tensor([0, 0, 1, 2, 1, 1])
</code></pre>
<p>the expected output would be three tensors:</p>
<pre><code class="lang-auto">torch.tensor([1])
torch.tensor([0, 2, 1])
torch.tensor([0, 1])
</code></pre>
</blockquote>
</aside>
<p>Take an <code>==</code> “outer product” of <code>samples</code> with <code>samples.unique()</code> (using<br>
broadcasting) to get a mask of the unique-value locations.  Use <code>where()</code><br>
to get the corresponding indices and use them to index into <code>samples</code>.<br>
Then split that result using the unique-value counts.</p>
<p>Thus:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.0
&gt;&gt;&gt;
&gt;&gt;&gt; samples = torch.tensor([50, 60, 40, 50, 50, 60])
&gt;&gt;&gt; labels = torch.tensor([0, 0, 1, 2, 1, 1])
&gt;&gt;&gt;
&gt;&gt;&gt; vals, cnts = samples.unique (return_counts = True)
&gt;&gt;&gt;
&gt;&gt;&gt; umask = samples.unsqueeze(0) == vals.unsqueeze (1)
&gt;&gt;&gt; umask   # just to check what umask looks like
tensor([[False, False,  True, False, False, False],
        [ True, False, False,  True,  True, False],
        [False,  True, False, False, False,  True]])
&gt;&gt;&gt;
&gt;&gt;&gt; inds = torch.where (umask)[1]
&gt;&gt;&gt; inds    # just to check what inds looks like
tensor([2, 0, 3, 4, 1, 5])
&gt;&gt;&gt;
&gt;&gt;&gt; # index into labels and split according to unique-value counts
&gt;&gt;&gt; labels[inds].split (cnts.tolist())
(tensor([1]), tensor([0, 2, 1]), tensor([0, 1]))
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

damiaanr: <aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="183270">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p><code>inds = torch.where (umask)[1]</code></p>
</blockquote>
</aside>
<p>Wow!!! This is clever! That was cool to see! <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>
<p>Thank you so much!</p>

damiaanr: <p>A quick addition:</p>
<p>In practice, if samples do actually not consist of numbers, but, e.g., stacked multi-dimensional tensors (let’s say every sample is <strong>torch.Size([26, 50])</strong> and the example <strong>samples</strong> tensor of tensor samples equals to <strong>torch.Size([6, 26, 50])</strong> in this case), then one can similarly use the <strong>==</strong> operator like Frank explains, but needs to apply <strong>.all(dim=n)</strong> for every dimension ‘belonging’ to the sample itself and not the indexing of the sample (in the example <strong>.all(dim=3).all(dim=2)</strong>).</p>
<p>However, the above is computationally way too expensive for real-world datasets, even if they are relatively small. Therefore, in the particular context of my question, a solution would be to pass an extra <strong>return_inverse=True</strong> parameter to the <strong>torch.unique()</strong> call (as in Frank’s code) and use the inverse counts as a replacement for <strong>samples</strong>. A light-weight replacement for <strong>vals</strong> can then simply be the range of the size of the first dimension of <strong>vals</strong> (i.e., <strong>torch.arange(0, vals.shape[0])</strong>. This should solve the memory problem for many cases. In one part of my particular situation, the samples tensor is 7,225,540 elements long, with 45,192 unique vals, so I have now implemented a hybrid setting to handle these extreme cases (determined by the availability of RAM on my machine <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"> ), just using the simple loop. Would be great if .eq() would yield booleans that take a bit instead of a byte, haha!</p>