Unexpected NaN gradient propagated when applying boolean conditions

username: rbregier
cooked: <p>Hi,</p>
<p>I am facing an unexpected autograd behavior when applying boolean conditions.</p>
<p>Below is a minimal example:</p>
<p>I evaluate the sinc function (equal to sin(x)/x for x!=0 and 1 otherwise), using ether the base definition of sinc or a Taylor series expansion close to 0 to avoid numerical issues.</p>
<pre><code class="lang-python">import torch

for value in (0.0, 1e-7, 1e-3):
    x = torch.as_tensor(value).requires_grad_(True)
    # Base definition valid for x != 0
    sinc_base = torch.sin(x)/x
    # Taylor series expansion valid for x close to 0
    sinc_taylor = 1.0 - x**2 / 6 
    # Switching between the two expressions depending on the value of x
    condition = torch.abs(x) &lt; 1e-6
    sinc = torch.where(condition, sinc_taylor, sinc_base)

    print(f"---- x={x.item()} ----")
    print(f"sinc_base: {sinc_base.item()} -- grad: {torch.autograd.grad(sinc_base, x, retain_graph=True)[0].item()}")
    print(f"sinc_taylor: {sinc_taylor.item()} -- grad: {torch.autograd.grad(sinc_taylor, x, retain_graph=True)[0].item()}")
    print(f"condition:", condition.item())
    print(f"sinc: {sinc.item()} -- grad: {torch.autograd.grad(sinc, x, retain_graph=True)[0].item()}")
    
    
# Output:
# ---- x=0.0 ----
# sinc_base: nan -- grad: nan
# sinc_taylor: 1.0 -- grad: -0.0
# condition: True
# sinc: 1.0 -- grad: nan
# ---- x=1.0000000116860974e-07 ----
# sinc_base: 1.0 -- grad: 0.0
# sinc_taylor: 1.0 -- grad: -3.33333360913457e-08
# condition: True
# sinc: 1.0 -- grad: -3.33333360913457e-08
# ---- x=0.0010000000474974513 ----
# sinc_base: 0.9999998807907104 -- grad: -0.0003662109375
# sinc_taylor: 0.9999998211860657 -- grad: -0.0003333333588670939
# condition: False
# sinc: 0.9999998807907104 -- grad: -0.0003662109375
</code></pre>
<p>The gradient of sinc should be finite everywhere. For x=0.0 however, the Taylor series sinc_taylor is used to estimate sinc, yet autograd does not return the gradient of sinc_taylor but NaN.</p>
<p>Do you think it is a bug, or do you have any workaround to suggest?</p>

username: rbregier
cooked: <p>Here is an even simpler example to illustrate the issue I am facing:</p>
<pre><code class="lang-python">import torch

x = torch.zeros(1, requires_grad=True)
x.squeeze().backward()
print(f"{x=} -- {x.grad=}")
# The gradient of x with respect to itself should be 1.
# We get the expected result:
# --&gt; x=tensor([0.], requires_grad=True) -- x.grad=tensor([1.])

x = torch.zeros(1, requires_grad=True)
other_value = 2 * x + 1
condition = torch.as_tensor(True)
y = torch.where(condition, x, other_value)
y.squeeze().backward()
print(f"{y=} -- {x.grad=}")
# We define y equal to x using a boolean condition.
# The gradient of y with respect to x should be equal to 1.
# We get the expected result:
# --&gt; y=tensor([0.], grad_fn=&lt;SWhereBackward0&gt;) -- x.grad=tensor([1.])

x = torch.zeros(1, requires_grad=True)
nan = 0.0 / x
condition = torch.as_tensor(True)
y = torch.where(condition, x, nan)
y.squeeze().backward()
print(f"{y=} -- {x.grad=}")
# We define y equal to x using a boolean condition.
# The gradient of y with respect to x should be equal to 1.
# We get nan:
# --&gt; y=tensor([0.], grad_fn=&lt;SWhereBackward0&gt;) -- x.grad=tensor([nan])
</code></pre>

username: InnovArul
cooked: <p>I could think of the following scenario:</p>
<p><code>torch.where(condition, x, y)</code> passes back <code>0</code> gradient to the elements that are not selected by the <code>condition</code>.</p>
<pre><code class="lang-auto">x = torch.zeros(1, requires_grad=True)
out = 0.0 / x
condition = torch.as_tensor(True)
y = torch.where(condition, x, out)
</code></pre>
<p>During backward, the gradient of <code>out</code> w.r.t <code>x</code> (<code>d(out) / dx</code>) would encounter <code>nan * 0</code>, which will result in <code>nan</code>.</p>

username: KFrank
cooked: <p>Hi rbregier!</p>
<aside class="quote no-group" data-username="rbregier" data-post="1" data-topic="153176" 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/rbregier/48/50191_2.png" class="avatar"> rbregier:</div>
<blockquote>
<p>yet autograd does not return the gradient of sinc_taylor but NaN.</p>
</blockquote>
</aside>
<p>Arul’s explanation is correct.  <code>torch.where()</code> backpropagates a gradient<br>
of <code>0 * nan = nan</code> (or maybe <code>0 * inf = nan</code>) through the “branch not<br>
taken.”</p>
<p>An example:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.10.2'
&gt;&gt;&gt; s1 = torch.zeros (4, requires_grad = True)
&gt;&gt;&gt; s2 = torch.zeros (4, requires_grad = True)
&gt;&gt;&gt; t1 = torch.tensor ([0.5, 0.5, 0.0, 0.0])
&gt;&gt;&gt; t2 = torch.tensor ([0.5, 0.5, 1.e-7, 1.e-7])
&gt;&gt;&gt; torch.where (t1 &gt; 1.e-6, s1 / t1, s1).sum().backward()
&gt;&gt;&gt; torch.where (t2 &gt; 1.e-6, s2 / t2, s2).sum().backward()
&gt;&gt;&gt; s1.grad
tensor([2., 2., nan, nan])
&gt;&gt;&gt; s2.grad
tensor([2., 2., 1., 1.])   # no gradients of 1.e7
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>Do you think it is a bug</p>
</blockquote>
</aside>
<p>It’s certainly a known issue, but I don’t think the devs consider it a bug.<br>
Apparently, it’s caused deep inside how autograd works with <code>where()</code><br>
and would be difficult to fix.</p>
<p>This github issues gives some explanation:</p>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/23395">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/23395" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/23395" target="_blank" rel="noopener nofollow ugc">Incorrect gradients for torch.where when one of the target tensors contains inf/nan</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2019-07-25" data-time="19:40:09" data-timezone="UTC">07:40PM - 25 Jul 19 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2019-07-25" data-time="20:00:47" data-timezone="UTC">08:00PM - 25 Jul 19 UTC</span>
        </div>

      <div class="user">
        <a href="https://github.com/egrefen" target="_blank" rel="noopener nofollow ugc">
          <img alt="egrefen" src="https://avatars.githubusercontent.com/u/1145770?v=4" class="onebox-avatar-inline" width="20" height="20">
          egrefen
        </a>
      </div>
    </div>

    <div class="labels">
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🐛 Bug

The `grad_fn` of `torch.where` returns the gradients of the wrong ar<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">gument, rather than of the selected tensor, if the other tensor's gradients have infs or nans.

## To Reproduce

Run this code:
```python
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("`torch.where` communicates gradients correctly:", torch.equal(expected_dw_dx, dwdx))
```

## Expected behavior

I would expect `expected_dw_dx == dwdx` in the example above.

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
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
```

PyTorch version: 1.1.0
Is debug build: No
CUDA used to build PyTorch: 9.0.176

OS: Ubuntu 18.04.1 LTS
GCC version: (Ubuntu 7.3.0-27ubuntu1~18.04) 7.3.0
CMake version: version 3.10.2

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: 9.0.176
GPU models and configuration:
GPU 0: Quadro GP100
GPU 1: Quadro GP100

Nvidia driver version: 410.79
cuDNN version: Could not collect

Versions of relevant libraries:
[pip] numpy==1.16.4
[pip] torch==1.1.0
[pip] torchvision==0.3.0
[conda] torch                     1.1.0                     &lt;pip&gt;
[conda] torchvision               0.3.0                     &lt;pip&gt;</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<aside class="quote no-group">
<blockquote>
<p>or do you have any workaround to suggest?</p>
</blockquote>
</aside>
<p>My approach is to get rid of the <code>nan</code>s.  You can safely feed an <em>incorrect</em><br>
value to the “zero” branch of <code>torch.where()</code>, as long as it’s not <code>nan</code><br>
(or <code>inf</code>, etc.).</p>
<p>In your specific case, I would take advantage of the fact that <code>sinc()</code> is an<br>
<em>even function</em> (<code>sinc (-x) = sinc (x)</code>), and <code>clamp()</code> the denominator<br>
away from zero:</p>
<pre data-code-wrap="python"><code class="lang-python">sinc_base = torch.sin (x.abs()) / x.abs().clamp (1.e-7)
</code></pre>
<p>So two things happen:  For <code>x.abs() &lt; 1.e-7</code>, <code>sinc_base</code> will be an<br>
<em>incorrect</em> value, but it won’t be <code>nan</code>.  However, for <code>x.abs (x) &lt; 1.e-6</code>,<br>
<code>torch.where()</code> will switch you over to <code>sinc_taylor</code>, so you will never<br>
see the incorrect <code>sinc_base</code> values.</p>
<p>Then, for gradients, when <code>abs (x) &lt; 1.e-6</code>, <code>torch.where()</code> will (in part)<br>
backpropagate <code>0 * sinc_base_gradient</code>.  Although <code>sinc_base_gradient</code><br>
will be incorrect for <code>abs (x) &lt; 1.e-7</code>, it won’t be <code>nan</code>, so autograd will<br>
correctly backpropagate <code>0</code> (rather than <code>nan</code>) for this piece of the of the<br>
<code>abs (x) &lt; 1.e-6</code> branch.</p>
<p>Best.</p>
<p>K. Frank</p>

username: rbregier
cooked: <p>Thank you Arul and K. Frank for your replies. I would not have thought that pytorch devs considers this as the expected behavior.  I will probably just try to avoid non finite values as you suggested.</p>