<p dir="auto">I think backprop through <code class="notranslate">torch.where</code> is wrong in certain special cases.</p>
<h2 dir="auto">To Reproduce</h2>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="t = 0
x = torch.ones(()).requires_grad_()
y = t * (x / t)  # just an example; anything that produces nan's works
z = torch.where(x &gt;= t, x, y)
z.backward()

# the forward pass works fine (the `nan`'s in `y` do not affect z)
# NOTE: this is unlike a naive implement of where that does `cond * x + (1 - cond) * y`
print(z)
# tensor(1., grad_fn=&lt;SWhereBackward&gt;)

# but the backward pass backprops the `nan`'s from y into x, even though the y path is never taken in torch.where
print(x.grad)
# tensor(nan)"><pre class="notranslate"><span class="pl-s1">t</span> <span class="pl-c1">=</span> <span class="pl-c1">0</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-en">requires_grad_</span>()
<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">t</span> <span class="pl-c1">*</span> (<span class="pl-s1">x</span> <span class="pl-c1">/</span> <span class="pl-s1">t</span>)  <span class="pl-c"># just an example; anything that produces nan's works</span>
<span class="pl-s1">z</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">where</span>(<span class="pl-s1">x</span> <span class="pl-c1">&gt;=</span> <span class="pl-s1">t</span>, <span class="pl-s1">x</span>, <span class="pl-s1">y</span>)
<span class="pl-s1">z</span>.<span class="pl-en">backward</span>()

<span class="pl-c"># the forward pass works fine (the `nan`'s in `y` do not affect z)</span>
<span class="pl-c"># NOTE: this is unlike a naive implement of where that does `cond * x + (1 - cond) * y`</span>
<span class="pl-en">print</span>(<span class="pl-s1">z</span>)
<span class="pl-c"># tensor(1., grad_fn=&lt;SWhereBackward&gt;)</span>

<span class="pl-c"># but the backward pass backprops the `nan`'s from y into x, even though the y path is never taken in torch.where</span>
<span class="pl-en">print</span>(<span class="pl-s1">x</span>.<span class="pl-s1">grad</span>)
<span class="pl-c"># tensor(nan)</span></pre></div>
<h2 dir="auto">Expected behavior</h2>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="print(x.grad)
# tensor(1.)"><pre class="notranslate"><span class="pl-en">print</span>(<span class="pl-s1">x</span>.<span class="pl-s1">grad</span>)
<span class="pl-c"># tensor(1.)</span></pre></div>
<p dir="auto">This would be the correct gradient.</p>
<p dir="auto">In practice, this bug can easily happen if one runs the above code for different t (including 0) and assumes that the <code class="notranslate">nan</code>'s for <code class="notranslate">t = 0</code> should not matter because the <code class="notranslate">torch.where</code> always selects the first path that has no <code class="notranslate">nan</code>s (and the forward pass does handle it correctly).</p>
<h2 dir="auto">Environment</h2>
<ul dir="auto">
<li>PyTorch Version 1.4</li>
</ul>
<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">The <code class="notranslate">grad_fn</code> of <code class="notranslate">torch.where</code> returns the gradients of the wrong argument, rather than of the selected tensor, if the other tensor's gradients have infs or nans.</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Run this code:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="x = torch.tensor([16., 0.], requires_grad=True)
y = x/2  # tensor([8., 0.], grad_fn=&lt;DivBackward0&gt;)
z = x.sqrt() + 1  # tensor([5., 1.], grad_fn=&lt;SqrtBackward&gt;)

# Calculate dy/dx, dz/dx
dydx = torch.autograd.grad(y.sum(), x, retain_graph=True)[0]  # tensor([0.5000, 0.5000])
dzdx = torch.autograd.grad(z.sum(), x, retain_graph=True)[0]  # tensor([0.1250,    inf])

# Define w = [w0, w1] == [y0, z1]
w = torch.where(x == 0., y, z)  # tensor([5., 0.], grad_fn=&lt;SWhereBackward&gt;)
expected_dw_dx = torch.where(x == 0., dydx, dzdx)  # tensor([0.1250, 0.5000])
dwdx = torch.autograd.grad(w.sum(), x, retain_graph=True)[0]  # is actually tensor([0.1250, inf])
print(&quot;`torch.where` communicates gradients correctly:&quot;, torch.equal(expected_dw_dx, dwdx))"><pre class="notranslate"><span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">tensor</span>([<span class="pl-c1">16.</span>, <span class="pl-c1">0.</span>], <span class="pl-s1">requires_grad</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)
<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span><span class="pl-c1">/</span><span class="pl-c1">2</span>  <span class="pl-c"># tensor([8., 0.], grad_fn=&lt;DivBackward0&gt;)</span>
<span class="pl-s1">z</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span>.<span class="pl-en">sqrt</span>() <span class="pl-c1">+</span> <span class="pl-c1">1</span>  <span class="pl-c"># tensor([5., 1.], grad_fn=&lt;SqrtBackward&gt;)</span>

<span class="pl-c"># Calculate dy/dx, dz/dx</span>
<span class="pl-s1">dydx</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">autograd</span>.<span class="pl-en">grad</span>(<span class="pl-s1">y</span>.<span class="pl-en">sum</span>(), <span class="pl-s1">x</span>, <span class="pl-s1">retain_graph</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)[<span class="pl-c1">0</span>]  <span class="pl-c"># tensor([0.5000, 0.5000])</span>
<span class="pl-s1">dzdx</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">autograd</span>.<span class="pl-en">grad</span>(<span class="pl-s1">z</span>.<span class="pl-en">sum</span>(), <span class="pl-s1">x</span>, <span class="pl-s1">retain_graph</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)[<span class="pl-c1">0</span>]  <span class="pl-c"># tensor([0.1250,    inf])</span>

<span class="pl-c"># Define w = [w0, w1] == [y0, z1]</span>
<span class="pl-s1">w</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">where</span>(<span class="pl-s1">x</span> <span class="pl-c1">==</span> <span class="pl-c1">0.</span>, <span class="pl-s1">y</span>, <span class="pl-s1">z</span>)  <span class="pl-c"># tensor([5., 0.], grad_fn=&lt;SWhereBackward&gt;)</span>
<span class="pl-s1">expected_dw_dx</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">where</span>(<span class="pl-s1">x</span> <span class="pl-c1">==</span> <span class="pl-c1">0.</span>, <span class="pl-s1">dydx</span>, <span class="pl-s1">dzdx</span>)  <span class="pl-c"># tensor([0.1250, 0.5000])</span>
<span class="pl-s1">dwdx</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">autograd</span>.<span class="pl-en">grad</span>(<span class="pl-s1">w</span>.<span class="pl-en">sum</span>(), <span class="pl-s1">x</span>, <span class="pl-s1">retain_graph</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)[<span class="pl-c1">0</span>]  <span class="pl-c"># is actually tensor([0.1250, inf])</span>
<span class="pl-en">print</span>(<span class="pl-s">"`torch.where` communicates gradients correctly:"</span>, <span class="pl-s1">torch</span>.<span class="pl-en">equal</span>(<span class="pl-s1">expected_dw_dx</span>, <span class="pl-s1">dwdx</span>))</pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">I would expect <code class="notranslate">expected_dw_dx == dwdx</code> in the example above.</p>
<h2 dir="auto">Environment</h2>
<p dir="auto">Please copy and paste the output from our<br>
<a href="https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py" rel="nofollow">environment collection script</a><br>
(or fill out the checklist below manually).</p>
<p dir="auto">You can get the script and run it with:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py"><pre class="notranslate"><code class="notranslate">wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
</code></pre></div>
<p dir="auto">PyTorch version: 1.1.0<br>
Is debug build: No<br>
CUDA used to build PyTorch: 9.0.176</p>
<p dir="auto">OS: Ubuntu 18.04.1 LTS<br>
GCC version: (Ubuntu 7.3.0-27ubuntu1~18.04) 7.3.0<br>
CMake version: version 3.10.2</p>
<p dir="auto">Python version: 3.7<br>
Is CUDA available: Yes<br>
CUDA runtime version: 9.0.176<br>
GPU models and configuration:<br>
GPU 0: Quadro GP100<br>
GPU 1: Quadro GP100</p>
<p dir="auto">Nvidia driver version: 410.79<br>
cuDNN version: Could not collect</p>
<p dir="auto">Versions of relevant libraries:<br>
[pip] numpy==1.16.4<br>
[pip] torch==1.1.0<br>
[pip] torchvision==0.3.0<br>
[conda] torch                     1.1.0                     <br>
[conda] torchvision               0.3.0                     </p>