Interesting argmax differential

trick: <p>I have two code blocks with argmax in both.<br>
First,</p>
<pre><code class="lang-auto">x = torch.tensor([0.1, 0.2], dtype=torch.float32, requires_grad=True)
y = torch.argmax(x, dim=-1)
y = torch.unsqueeze(y, 0)
y = y.to(torch.float32)
y = nn.Linear(1, 1)(y)
y.backward()
</code></pre>
<p>And second,</p>
<pre><code class="lang-auto">x = torch.tensor([0.1, 0.2], dtype=torch.float32, requires_grad=True)
y = torch.argmax(x, dim=-1)
y.backward()
</code></pre>
<p>Interesting thing is, the first code can run successful, but the second failed. So why?</p>

KFrank: <p>Hi Trick!</p>
<aside class="quote no-group" data-username="trick" data-post="1" data-topic="119089" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/5fc32e/48.png" class="avatar"> trick:</div>
<blockquote>
<pre><code class="lang-auto">y = torch.argmax(x, dim=-1)
</code></pre>
</blockquote>
</aside>
<p>As a side note, at this point <code>y</code> is <code> long</code> so it can’t have<br>
<code>requires_grad = True</code>.  So you wouldn’t be able to run<br>
<code>.backward()</code> on it.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">y = y.to(torch.float32)
y = nn.Linear(1, 1)(y)
y.backward()
</code></pre>
</blockquote>
</aside>
<p>Here, you make <code>y</code> a <code>float</code> (but it still doesn’t have<br>
<code>requires_grad = True</code>).  Then when you run it through <code>Linear</code>,<br>
it gets multiplied my <code>Linear</code>’s <code>weight</code> which <em>does</em> have<br>
<code>requires_grad = True</code> (and gets <code>bias</code> added to it which also has<br>
<code>requires_grad = True</code>).  So the result has <code>requires_grad = True</code>,<br>
and you can run <code>.backward()</code> on it.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">x = torch.tensor([0.1, 0.2], dtype=torch.float32, requires_grad=True)
y = torch.argmax(x, dim=-1)
y.backward()
</code></pre>
</blockquote>
</aside>
<p>Here, the fact that <code>x</code> has <code>requires_grad = True</code> is irrelevant.  The<br>
result of <code>argmax()</code> is  <code>long</code> (which can’t have <code>requires_grad = True</code>),<br>
so you can’t run <code>.backward()</code> on it.</p>
<p>Best.</p>
<p>K. Frank</p>

trick: <p>Hi,KFrank,<br>
Thanks for your reply. I am still confused. In the first code, when <code>y.backward()</code> is called, I think that the backward operation will be run through the first two lines , i.e. <code>dy/dx</code> will be calculated, but <code>argmax</code> is not differentiable, so it should cause an exception. But this seems to be wrong by your explanations.<br>
According to your explanations, the backward operation will not be run through the first two lines, because y does not have requires_grad, right? So the question is, in Pytorch, the <code>dy/dx</code> will be set to ZERO like differential of ReLU at point zero, or simply the backward operation has been stopped at earlier, for instance at line 5?</p>

KFrank: <p>Hi Trick!</p>
<aside class="quote no-group" data-username="trick" data-post="3" data-topic="119089" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/5fc32e/48.png" class="avatar"> trick:</div>
<blockquote>
<p>In the first code, when <code>y.backward()</code> is called, I think that the backward operation will be run through the first two lines , i.e. <code>dy/dx</code> will be calculated</p>
</blockquote>
</aside>
<p>This does not happen because, as you note:</p>
<aside class="quote no-group">
<blockquote>
<p><code>argmax</code> is not differentiable,</p>
</blockquote>
</aside>
<p>so it “breaks the computation graph” and the <code>.backward()</code> processing<br>
does not flow back to the first two lines so no attempt to calculate <code>dy/dy</code><br>
is made.</p>
<aside class="quote no-group">
<blockquote>
<p>so it should cause an exception.</p>
</blockquote>
</aside>
<p>As an aside, if there were <em>no</em> path backward through the computation<br>
graph, <code>.backward()</code> would raise and exception.  However (see below),<br>
there is another path backwards to a leaf with <code>requires_grad = True</code>,<br>
so, in your example, no exception is raised.</p>
<aside class="quote no-group">
<blockquote>
<p>the backward operation will not be run through the first two lines, because y does not have requires_grad, right?</p>
</blockquote>
</aside>
<p>Correct.</p>
<aside class="quote no-group">
<blockquote>
<p>So the question is, in Pytorch, the <code>dy/dx</code> will be set to ZERO like differential of ReLU at point zero, or simply the backward operation has been stopped at earlier, for instance at line 5?</p>
</blockquote>
</aside>
<p>The part of the backward operation that is trying to flow back to <code>x</code> is<br>
stopped (although I would say that it is stopped at line 4).  So <code>dy/dx</code><br>
is not set to zero, but, rather, is never created.  More precisely, <code>x.grad</code><br>
is never created.  And, just to say it explicitly, this is different than<br>
creating <code>x.grad</code> and having it be zero.</p>
<p>Now, as for there being another backward path, please note that<br>
<code>Linear</code> contains tensors (<code>weight</code> and <code>bias</code>) that have<br>
<code>requires_grad = True</code>.  So, even though <code>.backward()</code> does not<br>
create and calculate a <code>.grad</code> for <code>x</code>, it does for the tensors in <code>Linear</code>.</p>
<p>(And because there are some <code>.grad</code>s that are being calculated when<br>
you run <code>.backward()</code>, no exception is raised.)</p>
<p>Here is a short example that illustrates this:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.7.1'
&gt;&gt;&gt; x = torch.tensor([0.1, 0.2], dtype=torch.float32, requires_grad=True)
&gt;&gt;&gt; y = torch.argmax(x, dim=-1)
&gt;&gt;&gt; y = torch.unsqueeze(y, 0)
&gt;&gt;&gt; y = y.to(torch.float32)
&gt;&gt;&gt; my_linear = torch.nn.Linear(1, 1)
&gt;&gt;&gt; y = my_linear (y)
&gt;&gt;&gt; x.requires_grad
True
&gt;&gt;&gt; x.grad
&gt;&gt;&gt; my_linear.weight.requires_grad
True
&gt;&gt;&gt; my_linear.weight.grad
&gt;&gt;&gt; y.backward()
&gt;&gt;&gt; x.grad
&gt;&gt;&gt; my_linear.weight.grad
tensor([[1.]])
</code></pre>
<p>Notice that <code>x.grad</code> never gets created, and that <code>my_linear.weight.grad</code><br>
only gets created (and calculated) when <code>y.backward()</code> is called.</p>
<p>Best.</p>
<p>K. Frank</p>

trick: <p>Thanks, K. Frank. Now I’ve hot it</p>