Torch.where weird behavior

i4ata: <p>I have the following expression:</p>
<pre><code class="lang-auto">import torch
torch.manual_seed(0)

x = torch.rand(3,7,7,5)
mask1 = x[:,:,:,4] &gt; .5 # &lt;- [3,7,7]
x1, x2 = torch.rand_like(x), torch.rand_like(x)
mask2 = torch.rand(len(x[mask1])) &gt; .5 # &lt;- [74]

print(x[mask1].size(), mask2.size(), x1[mask1].size(), x2[mask1].size())
x[mask1] = torch.where(mask2, x1[mask1], x2[mask1])
</code></pre>
<p>The output of the print statement is:</p>
<pre><code class="lang-auto">torch.Size([74, 5]) torch.Size([74]) torch.Size([74, 5]) torch.Size([74, 5])
</code></pre>
<p>It seems that my tensors are compatible with each other. However, the line with <code>torch.where</code> gives the following Runtime Error: <em>The size of tensor a (74) must match the size of tensor b (5) at non-singleton dimension 1</em>.</p>
<p>Does anyone know what is the problem here?</p>
<p>When I do the <code>torch.where</code> manually like this:</p>
<pre><code class="lang-auto">x[mask1][mask2] = x1[mask1][mask2]
x[mask1][~mask2] = x2[mask1][~mask2]
</code></pre>
<p>It seems to work fine and to me that is the same logic as with the <code>torch.where</code>. Can someone explain?</p>
<p>Edit: Fixed typo</p>

KFrank: <p>Hi i4!</p>
<aside class="quote no-group" data-username="i4ata" data-post="1" data-topic="198652" 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/i4ata/48/68386_2.png" class="avatar"> i4ata:</div>
<blockquote>
<pre><code class="lang-auto">x[mask1] = torch.where(mask2, x1[mask1], x2[mask2])
</code></pre>
</blockquote>
</aside>
<p>First, you have a typo in this line – you want <code>x2[mask1]</code> (rather than<br>
<code>x2[mask2]</code>).</p>
<p>Your core problem is that the arguments you pass to <code>where()</code> don’t<br>
have the same shapes, so <code>where()</code> tries to <a href="https://pytorch.org/docs/2.2/notes/broadcasting.html" rel="noopener nofollow ugc">broadcast</a> them, but the<br>
shapes aren’t broadcastable.</p>
<p>Specifically, the last dimension of <code>mask2</code> is <code>74</code>, while the last dimension<br>
of, for example, <code>x1[mask1]</code> is <code>5</code>.  Adding a trailing singleton dimension<br>
to <code>mask2</code> (and fixing the typo) should do what you want:</p>
<pre data-code-wrap="python"><code class="lang-python">x[mask1] = torch.where (mask2.unsqueeze (-1), x1[mask1], x2[mask1])
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>When I do the <code>torch.where</code> manually like this:</p>
<pre><code class="lang-auto">x[mask1][mask2] = x1[mask1][mask2]
x[mask1][~mask2] = x2[mask1][~mask2]
</code></pre>
</blockquote>
</aside>
<p>As an aside, this won’t do what you think it will – it doesn’t actually modify<br>
<code>x</code>.</p>
<p><code>x[mask1] = x1[mask1]</code> would modify <code>x</code>, but, perhaps counterintuitively,<br>
when you index twice in a row, as in <code>x[mask1][mask2]</code>, somewhere along<br>
the line pytorch creates a new tensor, and you can’t use <code>x[mask1][mask2]</code><br>
to assign back into <code>x</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

i4ata: <p>Hello,</p>
<p>Thank you so much for your reply <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>