Is there a Pytorch low-level function to efficiently combine two tensors using a 0-1 mask?

zfzhang: <p>For example,<br>
tensors <strong>A</strong> &amp; <strong>B</strong> are of the same shape, <strong>m</strong> is a 0-1 tensor mask of the same shape,<br>
I want <code>C = A * m + B * (1-m)</code>, but implemented efficiently in low-level C++ code.</p>
<p>Thank you!</p>

OrielBanne: <p>it should work the way you wrote it given sizes are the same</p>

zfzhang: <p>Thanks! I’m looking for a faster low-level operation.</p>

OrielBanne: <p>I think numpy is as fast as any low level language (pytorch as well)</p>

Hamster: <p>Hello,<br>
You’re indeed right that 2 multiplies + 1 add (+1 sub for <code>1-m</code>) + the intermediary memory allocations are completely unnecessary (and also look horrible) when one could just do a CUDA kernel where the thread just looks at the boolean in m[thread_idx] and chooses either A[thread_idx] or B[thread_idx], which should be much faster. Fortunately, there’s already a native torch function which does just that: <code>torch.where</code> (<a href="https://pytorch.org/docs/stable/generated/torch.where.html" class="inline-onebox" rel="noopener nofollow ugc">torch.where — PyTorch 1.10.0 documentation</a>)</p>

vadimkantorov: <p>Another similar function is <code>torch.lerp</code> if your mask is a float tensor</p>

zfzhang: <p>Thanks, that’s precise! This function seems to have been introduce since v0.4, how can i miss it <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"> ?</p>

zfzhang: <p>Thanks! This is also a good solution, but I can only mark one as solution.</p>

KFrank: <p>Hi Zfzhang!</p>
<aside class="quote no-group" data-username="zfzhang" data-post="8" data-topic="134237" 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/z/ecc23a/48.png" class="avatar"> zfzhang:</div>
<blockquote>
<p>This is also a good solution</p>
</blockquote>
</aside>
<p>Please note – not necessarily relevant to your use case – that<br>
neither your original mask expression nor the more-or-less<br>
equivalent <code>torch.lerp()</code> protects against singularities in the<br>
values (nor gradients).</p>
<p>And <code>torch.where()</code> – perhaps surprisingly – doesn’t protect<br>
against singularities in the gradients.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt; x = torch.arange (-0.9, 1.0, 0.2, requires_grad = True)
&gt;&gt;&gt; m = (x &gt; 0).float()
&gt;&gt;&gt; y = m * torch.sqrt (x) + (1 - m) * torch.sqrt (-x)
&gt;&gt;&gt; y
tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
       grad_fn=&lt;AddBackward0&gt;)
&gt;&gt;&gt; y.sum().backward()
&gt;&gt;&gt; x.grad
tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan])
&gt;&gt;&gt; x = torch.arange (-0.9, 1.0, 0.2, requires_grad = True)
&gt;&gt;&gt; y = torch.where (x &gt; 0, torch.sqrt (x), torch.sqrt (-x))
&gt;&gt;&gt; y
tensor([0.9486832619, 0.8366600275, 0.7071067691, 0.5477225780, 0.3162277639,
        0.3162277639, 0.5477225780, 0.7071067691, 0.8366600275, 0.9486832619],
       grad_fn=&lt;SWhereBackward&gt;)
&gt;&gt;&gt; y.sum().backward()
&gt;&gt;&gt; x.grad
tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan])
&gt;&gt;&gt; x = torch.arange (-0.9, 1.0, 0.2, requires_grad = True)
&gt;&gt;&gt; y = torch.sqrt (x.abs())
&gt;&gt;&gt; y
tensor([0.9486832619, 0.8366600275, 0.7071067691, 0.5477225780, 0.3162277639,
        0.3162277639, 0.5477225780, 0.7071067691, 0.8366600275, 0.9486832619],
       grad_fn=&lt;SqrtBackward&gt;)
&gt;&gt;&gt; y.sum().backward()
&gt;&gt;&gt; x.grad
tensor([-0.5270463228, -0.5976142883, -0.7071067691, -0.9128708839,
        -1.5811388493,  1.5811388493,  0.9128708839,  0.7071067691,
         0.5976142883,  0.5270463228])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

zfzhang: <p>Thanks! This is what i expected.</p>