Compare two tensors and return total number of matches and the matching values

username: M2M
cooked: <p>Hi all,</p>
<p>I have two same-size 3D tensors, each one is populated with 4 possible values (0,1,2,3)</p>
<p>I want to compare these two tensors element-wise and return the total number of matches where both tensors had 0 in the same index, and the total number of matches where both tensors had 1 in the same index etc until the matching value is 3.</p>
<p>Any efficient way of doing this? Thank you!</p>

username: ZimoNitrome
cooked: <p>Does it need to be differentiable?</p>
<p>Otherwise:</p>
<pre><code class="lang-python">a = torch.randint(0, 4, [2,3,4])
b = torch.randint(0, 4, [2,3,4])
print(a)
print(b)

print(torch.unique(a[a == b], return_counts=True))
</code></pre>
<p>output:</p>
<pre><code class="lang-python">tensor([[[2, 1, 2, 1],
         [0, 0, 0, 0],
         [0, 0, 0, 1]],

        [[1, 3, 2, 3],
         [3, 3, 0, 0],
         [3, 1, 3, 3]]])

tensor([[[0, 3, 1, 3],
         [2, 0, 3, 3],
         [2, 0, 3, 1]],

        [[2, 2, 2, 1],
         [0, 2, 3, 0],
         [1, 0, 0, 2]]])

(tensor([0, 1, 2]), tensor([3, 1, 1]))
</code></pre>

username: M2M
cooked: <p>Thank you so much. This is what I needed. Just a quick question, can you please tell me what does<br>
a[a == b] do? I have never seen it before.</p>
<p>Also, in your example, if there is no matching for 3, so it would not return a value for that?</p>
<p>Thanks again!</p>

username: ZimoNitrome
cooked: <aside class="quote no-group" data-username="M2M" data-post="3" data-topic="137223">
<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/m2m/48/33611_2.png" class="avatar"> M2M:</div>
<blockquote>
<p>you please tell me what does a[a == b] do</p>
</blockquote>
</aside>
<p>We can divide it into 2 steps:</p>
<p><code>a == b</code> returns a boolean tensor (a mask) where values are <code>True</code> if the both <code>a</code> and <code>b</code> has the same value. The good thing is that with PyTorch, this operation is performed element-wise. So it checks each item at each channel, column, row and performs this operation.</p>
<p>Once we have this mask, we can use it to select items from the original tensors. It doesn’t matter if we were to do <code>a[a == b]</code> or <code>b[a == b]</code> since we only extract the values where <code>a</code> and <code>b</code> are the same.</p>
<p>We can see how this works in practice:</p>
<pre data-code-wrap="python"><code class="lang-python">a = torch.randint(0, 4, [2,3,4])
b = torch.randint(0, 4, [2,3,4])
print(a)
print(b)

mask = a == b
print(mask)

extracted_values = a[mask]
print(extracted_values)
</code></pre>
<p>Output:</p>
<pre data-code-wrap="python"><code class="lang-python"># a
tensor([[[0, 1, 2, 3],
         [3, 3, 2, 2],
         [1, 0, 0, 0]],

        [[1, 2, 1, 1],
         [3, 2, 3, 3],
         [3, 2, 0, 3]]])

# b
tensor([[[3, 3, 0, 0],
         [1, 2, 0, 3],
         [3, 0, 3, 0]],

        [[1, 0, 1, 3],
         [3, 0, 1, 1],
         [2, 2, 2, 1]]])

# mask
tensor([[[False, False, False, False],
         [False, False, False, False],
         [False,  True, False,  True]],

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

# extracted_values
tensor([0, 0, 1, 1, 3, 2])
</code></pre>
<p>We then count how many times each value appear in <code>extracted_values</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">values, counts = torch.unique(extracted_values, return_counts=True)
print(values)
print(counts)
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">tensor([0, 1, 2, 3])
tensor([2, 2, 1, 1])
</code></pre>
<p>So we can see for these new random tensors, a and b has 2 '0’s in the same position, 2 '1’s in the same position, 1 ‘2’ in the same position, and 1 ‘3’ in the same position.</p>
<aside class="quote no-group" data-username="M2M" data-post="3" data-topic="137223">
<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/m2m/48/33611_2.png" class="avatar"> M2M:</div>
<blockquote>
<p>Also, in your example, if there is no matching for 3, so it would not return a value for that?</p>
</blockquote>
</aside>
<p>Exactly, that is how <code>torch.unique()</code> works.</p>

username: M2M
cooked: <p>Thank you so much. This is so helpful. a == b will create a mask and then we can pass this mask to extract values from tensors is super cool of Pytorch. Thanks again !!!</p>