<p dir="auto">Porting TH operators is essential for code simplicity and performance reasons.</p>
<p dir="auto">Porting guides and Q&amp;A are available in umbrella issue: <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="481740397" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/24507" data-hovercard-type="issue" data-hovercard-url="/pytorch/pytorch/issues/24507/hovercard" href="https://github.com/pytorch/pytorch/issues/24507">#24507</a></p>
<p dir="auto">Feel free to add <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/VitalyFedyunin/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/VitalyFedyunin">@VitalyFedyunin</a> as a reviewer to get a prioritized review.</p>
<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">This is a niche bug, but it might cause troubles in advanced users who like to use masking to filter out NaN losses. Simply put, when NaN losses are masked out using <code class="notranslate">masked_fill</code>, performing <code class="notranslate">backward</code> on the sum of the losses should produce valid gradients (assuming that the gradient graph is smooth everywhere except for the masked losses). When</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Steps to reproduce the behavior:</p>
<p dir="auto">Define the environment as follows. We will be backpropagating gradients to a very simple <code class="notranslate">nn.Linear</code> layer.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import torch
import torch.nn as nn
x = torch.ones(3, 4)
linear = nn.Linear(4, 2)
def has_err(x):
    return bool(((x != x) | (x == float(&quot;inf&quot;)) | (x == float(&quot;-inf&quot;))).any().item())"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">torch</span>
<span class="pl-k">import</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span> <span class="pl-k">as</span> <span class="pl-s1">nn</span>
<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">ones</span>(<span class="pl-c1">3</span>, <span class="pl-c1">4</span>)
<span class="pl-s1">linear</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Linear</span>(<span class="pl-c1">4</span>, <span class="pl-c1">2</span>)
<span class="pl-k">def</span> <span class="pl-en">has_err</span>(<span class="pl-s1">x</span>):
    <span class="pl-k">return</span> <span class="pl-en">bool</span>(((<span class="pl-s1">x</span> <span class="pl-c1">!=</span> <span class="pl-s1">x</span>) <span class="pl-c1">|</span> (<span class="pl-s1">x</span> <span class="pl-c1">==</span> <span class="pl-en">float</span>(<span class="pl-s">"inf"</span>)) <span class="pl-c1">|</span> (<span class="pl-s1">x</span> <span class="pl-c1">==</span> <span class="pl-en">float</span>(<span class="pl-s">"-inf"</span>))).<span class="pl-en">any</span>().<span class="pl-en">item</span>())</pre></div>
<p dir="auto">Performing a forward inference on the linear layer and propagating gradients from there works as intended if no <code class="notranslate">nan</code> is involved.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="linear(x).sum().backward()
print(linear.weight.grad)
# tensor([[3., 3., 3., 3.],
#         [3., 3., 3., 3.]])
assert not has_err(linear.weight.grad)"><pre class="notranslate"><span class="pl-en">linear</span>(<span class="pl-s1">x</span>).<span class="pl-en">sum</span>().<span class="pl-en">backward</span>()
<span class="pl-en">print</span>(<span class="pl-s1">linear</span>.<span class="pl-s1">weight</span>.<span class="pl-s1">grad</span>)
<span class="pl-c"># tensor([[3., 3., 3., 3.],</span>
<span class="pl-c">#         [3., 3., 3., 3.]])</span>
<span class="pl-k">assert</span> <span class="pl-c1">not</span> <span class="pl-en">has_err</span>(<span class="pl-s1">linear</span>.<span class="pl-s1">weight</span>.<span class="pl-s1">grad</span>)</pre></div>
<p dir="auto">Suppose that <code class="notranslate">x</code> has <code class="notranslate">nan</code> in one of its rows.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="x[2].fill_(float(&quot;nan&quot;))
print(linear(x))
# tensor([[ 0.0570, -1.0066],
#         [ 0.0570, -1.0066],
#         [    nan,     nan]], grad_fn=&lt;AddmmBackward&gt;)"><pre class="notranslate"><span class="pl-s1">x</span>[<span class="pl-c1">2</span>].<span class="pl-en">fill_</span>(<span class="pl-en">float</span>(<span class="pl-s">"nan"</span>))
<span class="pl-en">print</span>(<span class="pl-en">linear</span>(<span class="pl-s1">x</span>))
<span class="pl-c"># tensor([[ 0.0570, -1.0066],</span>
<span class="pl-c">#         [ 0.0570, -1.0066],</span>
<span class="pl-c">#         [    nan,     nan]], grad_fn=&lt;AddmmBackward&gt;)</span></pre></div>
<p dir="auto">Performing backwards computation on the sum of the matrix above should<br>
produce <code class="notranslate">nan</code> gradients. This is expected.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="linear.zero_grad()
linear(x).sum().backward()
print(linear.weight.grad)
# tensor([[nan, nan, nan, nan],
#         [nan, nan, nan, nan]])
assert has_err(linear.weight.grad)"><pre class="notranslate"><span class="pl-s1">linear</span>.<span class="pl-en">zero_grad</span>()
<span class="pl-en">linear</span>(<span class="pl-s1">x</span>).<span class="pl-en">sum</span>().<span class="pl-en">backward</span>()
<span class="pl-en">print</span>(<span class="pl-s1">linear</span>.<span class="pl-s1">weight</span>.<span class="pl-s1">grad</span>)
<span class="pl-c"># tensor([[nan, nan, nan, nan],</span>
<span class="pl-c">#         [nan, nan, nan, nan]])</span>
<span class="pl-k">assert</span> <span class="pl-en">has_err</span>(<span class="pl-s1">linear</span>.<span class="pl-s1">weight</span>.<span class="pl-s1">grad</span>)</pre></div>
<p dir="auto">However, using <code class="notranslate">masked_fill</code> or  index slicing on the leaf nodes to mask out problematic gradient graphs should not.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="linear.zero_grad()
linear(x)[:2].sum().backward()
assert not has_err(linear.weight.grad)  # AssertionError raised"><pre class="notranslate"><span class="pl-s1">linear</span>.<span class="pl-en">zero_grad</span>()
<span class="pl-en">linear</span>(<span class="pl-s1">x</span>)[:<span class="pl-c1">2</span>].<span class="pl-en">sum</span>().<span class="pl-en">backward</span>()
<span class="pl-k">assert</span> <span class="pl-c1">not</span> <span class="pl-en">has_err</span>(<span class="pl-s1">linear</span>.<span class="pl-s1">weight</span>.<span class="pl-s1">grad</span>)  <span class="pl-c"># AssertionError raised</span></pre></div>
<p dir="auto">We can check whether masking out final losses works in cases where no <code class="notranslate">nan</code> is involved.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="x = torch.ones(3, 4)
x[2].fill_(0.5)
linear.zero_grad()
linear(x).sum().backward()
print(linear.weight.grad)
# without masking:
# tensor([[2.5000, 2.5000, 2.5000, 2.5000],
#         [2.5000, 2.5000, 2.5000, 2.5000]])
linear.zero_grad()
linear(x)[2:].sum().backward()
print(linear.weight.grad)
# with masking:
# tensor([[2., 2., 2., 2.],
#         [2., 2., 2., 2.]])"><pre class="notranslate"><span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">ones</span>(<span class="pl-c1">3</span>, <span class="pl-c1">4</span>)
<span class="pl-s1">x</span>[<span class="pl-c1">2</span>].<span class="pl-en">fill_</span>(<span class="pl-c1">0.5</span>)
<span class="pl-s1">linear</span>.<span class="pl-en">zero_grad</span>()
<span class="pl-en">linear</span>(<span class="pl-s1">x</span>).<span class="pl-en">sum</span>().<span class="pl-en">backward</span>()
<span class="pl-en">print</span>(<span class="pl-s1">linear</span>.<span class="pl-s1">weight</span>.<span class="pl-s1">grad</span>)
<span class="pl-c"># without masking:</span>
<span class="pl-c"># tensor([[2.5000, 2.5000, 2.5000, 2.5000],</span>
<span class="pl-c">#         [2.5000, 2.5000, 2.5000, 2.5000]])</span>
<span class="pl-s1">linear</span>.<span class="pl-en">zero_grad</span>()
<span class="pl-en">linear</span>(<span class="pl-s1">x</span>)[<span class="pl-c1">2</span>:].<span class="pl-en">sum</span>().<span class="pl-en">backward</span>()
<span class="pl-en">print</span>(<span class="pl-s1">linear</span>.<span class="pl-s1">weight</span>.<span class="pl-s1">grad</span>)
<span class="pl-c"># with masking:</span>
<span class="pl-c"># tensor([[2., 2., 2., 2.],</span>
<span class="pl-c">#         [2., 2., 2., 2.]])</span></pre></div>
<p dir="auto">The accumulated gradient in the latter case is less than the former because index slicing (<code class="notranslate">[2:]</code>) on the computation graph prevented the third row in <code class="notranslate">x</code> from affecting the gradient computation.</p>
<h2 dir="auto">Environment</h2>
<p dir="auto">Collecting environment information...<br>
PyTorch version: 1.0.0<br>
Is debug build: No<br>
CUDA used to build PyTorch: None</p>
<p dir="auto">OS: Mac OSX 10.14.5<br>
GCC version: Could not collect<br>
CMake version: Could not collect</p>
<p dir="auto">Python version: 3.7<br>
Is CUDA available: No<br>
CUDA runtime version: No CUDA<br>
GPU models and configuration: No CUDA<br>
Nvidia driver version: No CUDA<br>
cuDNN version: No CUDA</p>
<p dir="auto">Versions of relevant libraries:<br>
[pip] numpy==1.16.3<br>
[pip] torch==1.0.0<br>
[conda] torch                     1.0.0                     </p>
<h2 dir="auto">Additional context</h2>
<p dir="auto">None.</p>