Gradients of torch.where

Nicki: <p>Hello,</p>
<p>I am trying to calculate gradients of a function that uses torch.where, however it results in unexpected gradients. I basically use it to choose between some real case, complex case and limit case where some of the cases will have a Nan gradient for some specific input. For simplicity consider the following example:</p>
<pre><code>def f1(x):
    return 0/x

def f2(x):
    return x

def g(x):
    r1 = f1(x)
    r2 = f2(x)
    return torch.where(x == 0.0, r2, r1)

x = torch.autograd.Variable(torch.zeros(1), requires_grad=True)
print(torch.autograd.grad(f(x), x))
&gt;&gt; (tensor([nan]),)
print(torch.autograd.grad(f2(x), x))
&gt;&gt; (tensor([1.]),)
</code></pre>
<p>I would expect that the gradient of f(x) to be the same as f2(x) but they are clearly different. This problem only happens when one of the functions has a Nan gradient i.e. if I change f1 to something like x**2 I get the correct result. I do not know how gradients are propagated through the torch.where function, but if it is something like</p>
<pre><code>torch.where(condition, tensor1, tensor2) = condition * tensor1 + (1-condition) * tensor2
</code></pre>
<p>then it make sense that the resulting gradient is Nan, even if condition = 0 and only tensor2 has a Nan gradient.</p>
<p>Is this the intended behavior of torch.where or is there missing some Nan checking missing? If it is intended, it there a way to get around it?</p>
<p>Thank you</p>

tom: <p>While it is a bit confusing at first sight, there is little that can be done here, and in particular, <code>where</code> is working as expected.<br>
When you instrument this as</p>
<pre><code class="lang-python">def f1(x):
    return 0/x

def f2(x):
    return x.clone()

x = torch.zeros(1, requires_grad=True)
r1 = f1(x)
r1.retain_grad()
r2 = f2(x)
r2.retain_grad()
c = torch.where(x == 0, r2, r1)
c.backward()
print(r1.grad, r2.grad, x.grad)
</code></pre>
<p>Here you see that r1.grad is 0, r2.grad is 1 and x.grad is NaN.</p>
<p>The gradient x.grad is NaN because even when you feed a grad of 0 to the backward of f1, you’ll get NaN as the gradient of f1’s input.<br>
The gradient of x is the sum of the input gradients of f1 and f2, i.e. NaN + 1 = NaN.</p>
<p>Best regards</p>
<p>Thomas</p>

tom-programming: <p>It seems connected to a thing I noted a while ago: <a href="https://discuss.pytorch.org/t/bug-or-feature-nans-influence-other-variables-in-backprop/17049/3">Bug or feature? NaNs influence other variables in backprop</a><br>
In that case there was a 0* inside the code of PyTorch that did 0*NaN and obtained NaN. The workaround for me has been to substitute NaNs with zeros.</p>
<p>Anyway, even if it’s expected behavior, I don’t feel the PyTorch implementation is the most reasonable</p>

AlbertZeyer: <aside class="quote no-group" data-username="tom" data-post="2" data-topic="26835">
<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/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>The gradient of x is the sum of the input gradients of f1 and f2</p>
</blockquote>
</aside>
<p>No, this is wrong. (Considere <code>torch.where(..., r1, r1)</code>. With your argument, the grad should double, but clearly this is just the identity function.)</p>

tom: <p>Your example does not change the fact that the backward of “use the same tensor several times” is the addition of the gradients of the various uses.</p>
<p>In your example, conceptually, the gradient of the loss w.r.t. output of <code>where</code> is split between the second argument where the condition in the first argument is <code>True</code>/nonzero (which goes to the first r1 use) and the third argument (which goes to the second r1 use). Now because you put in the same tensor into the two arguments, these two uses will be added to get the identity, but they were masked copies of the gradient w.r.t. the output of where.</p>
<p>Just because you could simplify the expression does not mean that the autograd mechanics change for the operators you give. Of course, if you run this through a compiler that optimizes the case, the gradient computation might also change.</p>

AlbertZeyer: <p>So you say that the sum is w.r.t. to the masked tensor?</p>
<p>What I expect for the grad: For some given:</p>
<pre data-code-wrap="python"><code class="lang-python">c = where(mask, a, b)
</code></pre>
<p>For the grad of the loss w.r.t. <code>a</code> and <code>b</code>, specified as <code>da</code>/<code>db</code>, in dependence of the grad of the loss w.r.t. <code>c</code>, specified as <code>dc</code>:</p>
<pre><code class="lang-auto">da = where(mask, dc, 0.)
db = where(~mask, dc, 0.)
</code></pre>
<p>For the given earlier example with <code>mask = x == 0</code> (<code>== True</code>), <code>a=r2</code>, <code>b=r1</code>, it means:</p>
<pre data-code-wrap="python"><code class="lang-python">da = dc = 1
db = 0
</code></pre>
<p>I think you say that:</p>
<pre data-code-wrap="python"><code class="lang-python">dx = f2.backward(da) + f1.backward(db)
</code></pre>
<p>Ah, I think I see what you mean now. This will give:</p>
<pre data-code-wrap="python"><code class="lang-python">dx = 1. * da + 0.0 * log(x) * db
   = 1. + 0. * NaN
   = NaN
</code></pre>
<p>But speaking mathematically, this computation is wrong. The grad of <code>c</code> w.r.t. <code>x</code> is clearly 1. With <code>db = 0</code>, it does not matter that the grad of <code>r1</code> w.r.t. <code>x</code> is undefined (NaN).</p>
<p>I wonder if autograd should maybe fix this? The application of the chain rule is wrong: when <code>grad_output</code> is 0 somewhere, <code>grad_input</code> must also be 0. Currently it does not do that.</p>

tom: <aside class="quote no-group" data-username="AlbertZeyer" data-post="6" data-topic="26835">
<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/albertzeyer/48/45362_2.png" class="avatar"> AlbertZeyer:</div>
<blockquote>
<p>But speaking mathematically, this computation is wrong.</p>
</blockquote>
</aside>
<p>As a mathematician, I would say that the computation is mathematically defined as being IEEE 754 operations, including NaN trumping everything else.</p>
<p>But so leaving aside my reservations about going around and saying “you’re wrong” at people trying to help understand why things are the way they are, it seems that you are trying to start a conversation about the autograd behaviour involved: People have suggested using “masking zeros” for gradients before, it has not gone anywhere because a departure from IEEE semantics is quite expensive. Also, it is much less universally useful than it looks because not all operations have a pointwise backwards and the conventional wisdom is to avoid generating NaNs rather than try to deal with them.</p>
<p><a href="https://pytorch.org/docs/stable/masked.html" rel="noopener nofollow ugc">Masked tensors</a> don’t quite do what you want (they don’t add) but maybe you could define your own tensor subclass that takes inspiration there to achieve things.</p>

AlbertZeyer: <p>Thanks for the clarification. I think I understand the reasoning, and I came to the same conclusion, that currently it is best to avoid generating NaNs.</p>
<p>Although I still wonder whether there is a better solution, which is still efficient.</p>
<p>I’m not sure that I agree with the terminology. When I say “wrong”, I mean that the algorithm (the autograd implementation) does not correctly calculates the real gradient.</p>
<p>Or maybe you refer to “0 * Nan = Nan”? That’s not what I mean when I say it’s wrong. This is totally fine: As you say, this is just what IEEE defines.</p>
<p>I refer to what autograd calculates. It does not calculate the correct gradient. It sometimes (as in this example) calculates NaN even though the gradient is well defined. My assumption was always that autograd always calculates the correct gradient when it is well defined. But this assumption doesn’t seem to be the case because of the behavior as shown here.</p>