Torch backward PowBackward0 causes nan gradient where it shouldn't

Paolo_R: <p>I have a pytorch tensor with NaN inside, when I calculate the loss function using a simple MSE Loss the gradient becomes NaN even if I mask out the NaN values.</p>
<p>Weirdly this happens only when the mask is applyied after calculating the loss and only when the loss has a pow operation inside. The various cases follow</p>
<pre><code class="lang-auto">import torch
torch.autograd.set_detect_anomaly(True)

x = torch.rand(10, 10) 
y = torch.rand(10, 10)
w = torch.rand(10, 10, requires_grad=True)
y[y &gt; 0.5] = torch.nan


o = w @ x
l = (y - o)**2
l = l[~y.isnan()]

try:
    l.mean().backward(retain_graph=True)
except RuntimeError:
    print('(y-o)**2 caused nan gradient')

l = (y - o)
l = l[~y.isnan()]

try:
    l.mean().backward(retain_graph=True)
except RuntimeError():
    pass
else:
    print('y-o does not cause nan gradient')

l = (y[~y.isnan()] - o[~y.isnan()])**2
l.mean().backward()
print('masking before pow does not propagate nan gradient')
</code></pre>
<p>What makes NaN gradients propagate when passing through the backward pass of the pow function?</p>

KFrank: <p>Hi Paolo!</p>
<aside class="quote no-group" data-username="Paolo_R" data-post="1" data-topic="209647" 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/paolo_r/48/72779_2.png" class="avatar"> Paolo_R:</div>
<blockquote>
<pre><code class="lang-auto">l = (y - o)**2
</code></pre>
</blockquote>
</aside>
<p>Here, <code>l</code> contains <code>nan</code>s (because <code>y</code> contains <code>nan</code>s).</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">l = l[~y.isnan()]
</code></pre>
</blockquote>
</aside>
<p>Here you create a new tensor that you also call <code>l</code>.  This new<br>
tensor does not contain <code>nan</code>s.  But the new <code>l</code> depends on the<br>
old <code>l</code> which is saved in the computation graph and still contains<br>
<code>nan</code>s.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">  l.mean().backward(retain_graph=True)
</code></pre>
</blockquote>
</aside>
<p>Because both the new <code>l</code> and the old <code>l</code> depend on the leaf<br>
variable <code>w</code> that carries <code>requires_grad = True</code>, you have to<br>
backpropagate through the old <code>l</code>.  The gradient of the old <code>l</code><br>
with respect to <code>y</code> is <code>2 * y</code> which contains <code>nan</code>s for those<br>
elements of <code>y</code> that are <code>nan</code>. So you hit the error.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">l = (y - o)
l = l[~y.isnan()]
...
    l.mean().backward(retain_graph=True)
</code></pre>
</blockquote>
</aside>
<p>Again, you backpropagate through new <code>l</code>, old <code>l</code>, and <code>y</code>.  But<br>
now the gradient of old <code>l</code> with respect to <code>y</code> is just <code>1</code>, so the<br>
<code>nan</code>s in <code>y</code> don’t enter into the backpropagation.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">l = (y[~y.isnan()] - o[~y.isnan()])**2
l.mean().backward()
</code></pre>
</blockquote>
</aside>
<p>This time <code>y[~y.isnan()]</code> is a new tensor (that you haven’t given<br>
a name to) that does not contain <code>nan</code>s.  Likewise, <code>o[~y.isnan()]</code><br>
contains no <code>nan</code>s.The gradient of <code>l</code> with respect to <code>y[~y.isnan()]</code><br>
is <code>2 * y[~y.isnan()]</code> which doesn’t contain <code>nan</code>s, so there is no<br>
error.</p>
<p>Best.</p>
<p>K. Frank</p>