Applying a function on batched tensor

se7en.dreams: <p>I have the following tensors:</p>
<pre><code class="lang-python"># 2 x 5 x 3
a = torch.tensor([[[1, 3, 2], [7, 9, 8], [13, 15, 14], [19, 21, 20], [25, 27, 26]], [[31, 33, 32], [37, 39, 38], [43, 45, 46], [49, 51, 50], [55, 57, 56]]])
b = torch.tensor([[[19, 21, 20], [7, 9, 8], [13, 15, 14], [1, 3, 2], [25, 27, 26]], [[55, 57, 56], [31, 33, 32], [37, 39, 38], [43, 45, 44], [49, 51, 50]]])
</code></pre>
<p>I’d like to obtain the following:</p>
<pre><code class="lang-python"># 2 x 5 x 5
c = torch.tensor([[[0, 0, 0, 1, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 0, 0, 0, 1]], [[0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [0, 0, 0, 0, 1], [1, 0, 0, 0, 0]]])
</code></pre>
<p>where, the values of <code>c</code> are obtained using the function:</p>
<pre><code class="lang-python"># compare each value of a with each value of b. 
# If they are equal, then set the value at those indices as 1 in c, otherwise set it to 0. 
# For instance, a's first value [1, 3, 2] is compared with b's first value [19, 21, 20]. Since, they are not equal, I will set c's first value as 0. 
def test(x, y):
    if torch.all(x==y):
        return 1
    else:
        return 0
</code></pre>
<p>How do I do this without looping over <code>a</code> and <code>b</code>?</p>

KFrank: <p>Hi Seven!</p>
<aside class="quote no-group" data-username="se7en.dreams" data-post="1" data-topic="157955" data-full="true">
<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/s/a4c791/48.png" class="avatar"> se7en.dreams:</div>
<blockquote>
<p>I’d like to obtain the following:</p>
<pre data-code-wrap="python"><code class="lang-python"># 2 x 5 x 5
c = torch.tensor([[[0, 0, 0, 1, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 0, 0, 0, 1]], [[0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [0, 0, 0, 0, 1], [1, 0, 0, 0, 0]]])
</code></pre>
</blockquote>
</aside>
<p>I don’t understand your use case.  Among other things, the second row<br>
of your tensor, <code>c</code>, seems to indicate five matches between the second<br>
rows of <code>a</code> and <code>b</code>, but I only see four.</p>
<p>Could you post a complete, runnable script that illustrates your use case<br>
by computing your desired output?</p>
<aside class="quote no-group">
<blockquote>
<p>How do I do this without looping over <code>a</code> and <code>b</code>?</p>
</blockquote>
</aside>
<p>This approach – which may or may not be applicable to your specific<br>
use case – illustrates a loop-free way to compare the rows of <code>a</code> and <code>b</code>:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.12.0'
&gt;&gt;&gt; a = torch.tensor([[[1, 3, 2], [7, 9, 8], [13, 15, 14], [19, 21, 20], [25, 27, 26]], [[31, 33, 32], [37, 39, 38], [43, 45, 46], [49, 51, 50], [55, 57, 56]]])
&gt;&gt;&gt; b = torch.tensor([[[19, 21, 20], [7, 9, 8], [13, 15, 14], [1, 3, 2], [25, 27, 26]], [[55, 57, 56], [31, 33, 32], [37, 39, 38], [43, 45, 44], [49, 51, 50]]])
&gt;&gt;&gt; c = torch.tensor([[[0, 0, 0, 1, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 0, 0, 0, 1]], [[0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [0, 0, 0, 0, 1], [1, 0, 0, 0, 0]]])
&gt;&gt;&gt; cB = (a.unsqueeze (1) == b.unsqueeze (2)).prod (dim = 3)
&gt;&gt;&gt; torch.equal (c[0], cB[0])
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

se7en.dreams: <p>Thanks for your response, K. Frank. So, here is what I’m after (See <code>R</code> in image):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/a/1a14feb5497888802e2c2f657f04fefd8b8a26a8.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/1a14feb5497888802e2c2f657f04fefd8b8a26a8" title="IMG_2C2AA7E5EB4D-1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/a/1a14feb5497888802e2c2f657f04fefd8b8a26a8_2_666x500.jpeg" alt="IMG_2C2AA7E5EB4D-1" data-base62-sha1="3IJox7nRJ4d9qc5v7MID1wVV1FK" width="666" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/a/1a14feb5497888802e2c2f657f04fefd8b8a26a8_2_666x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/1/a/1a14feb5497888802e2c2f657f04fefd8b8a26a8_2_999x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/a/1a14feb5497888802e2c2f657f04fefd8b8a26a8_2_1332x1000.jpeg 2x" data-dominant-color="E8E8E9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">IMG_2C2AA7E5EB4D-1</span><span class="informations">1656×1242 157 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<pre data-code-wrap="python"><code class="lang-python">o = torch.tensor([[[1, 3, 2], [7, 9, 8], [13, 15, 14], [19, 21, 20], [25, 27, 26]], [[31, 33, 32], [37, 39, 38], [43, 45, 44], [49, 51, 50], [55, 57, 56]]])
p = torch.tensor([[[19, 21, 20], [7, 9, 8], [13, 15, 14], [1, 3, 2], [25, 27, 26]], [[55, 57, 56], [31, 33, 32], [37, 39, 38], [43, 45, 44], [49, 51, 50]]])

# this is O' in image
o_prime = torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5], [0.6, 0.7, 0.8, 0.9, 0.11]])

# this is P' in image
p_prime = torch.tensor([[1.1, 1.2, 1.3, 1.4, 1.5], [1.6, 1.7, 1.8, 1.9, 1.11]])

# this is R (this is what I need)
r = torch.tensor([[[0, 0, 0, 6.1, 0], [0, 24.2, 0, 0, 0], [0, 0, 42.3, 0, 0], [60.4, 0, 0, 0, 0], [0, 0, 0, 0, 78.5]], [[0, 96.6, 0, 0, 0], [0, 0, 114.7, 0, 0], [0, 0, 0, 132.8, 0], [0, 0, 0, 0, 150.9], [168.11, 0, 0, 0, 0]]])
</code></pre>
<p>Let me know if you need more clarification.</p>