Incorrect gradient calculation with torch.where and NaNs

username: Aeryan
cooked: <p>I am implementing a truncated gaussian distribution in Pytorch (<a href="https://github.com/TheAeryan/stable-truncated-gaussian" rel="noopener nofollow ugc">https://github.com/TheAeryan/stable-truncated-gaussian</a>). In order to do that in a numerically stable manner, I need conditional code (<em>if-then-else</em>) for calculating the mean, variance and log-probs in one way or another, depending on the specific input values. Here is an example:</p>
<pre><code class="lang-auto">def _F_1(x, y):
		if torch.abs(x) &gt; torch.abs(y):
			out = SequentialTruncatedGaussian._F_1(y, x)
		elif torch.abs(x - y) &lt; 1e-7:
			out = SequentialTruncatedGaussian._P_1(x, y-x)
		elif x &lt;= 0 and y &lt;= 0:
			delt = SequentialTruncatedGaussian._delta(x, y)
			out = (1 - delt) / (delt*erfcx(-y) - erfcx(-x))  
		elif x &gt;= 0 and y &gt;= 0:
			delt = SequentialTruncatedGaussian._delta(x, y)
			out = (1 - delt) / (erfcx(x) - delt*erfcx(y))
		else:
			delt = SequentialTruncatedGaussian._delta(x, y)
			out = ((1-delt)*torch.exp(-x**2)) / (erf(y)-erf(x))

		return out
</code></pre>
<p>I am now adapting my code to make it parallel, so I need a way to <em>insert</em> these <em>if-then-else</em> conditions into Pytorch. What I do is that I compute the <em>out</em> values associated with each condition and, then, I use <em>torch.where</em> to select the correct value according to its associated condition. This is how I have implemented the <em>_F_1</em> function above:</p>
<pre><code class="lang-auto">def _F_1(x_, y_):
		# All values in tensor @x must be smaller than values in tensor @y
		x = torch.where(torch.abs(x_)&lt;=torch.abs(y_),x_,y_)
		y = torch.where(torch.abs(y_)&gt;=torch.abs(x_),y_,x_)

		# Values
		delt = NewParallelTruncatedGaussian._delta(x, y)
		one_minus_delt = 1 - delt
		out1 = NewParallelTruncatedGaussian._P_1(x, y-x)
		out2 = one_minus_delt / (delt*erfcx(-y) - erfcx(-x) + EPS)
		out3 = one_minus_delt / (erfcx(x) - delt*erfcx(y) + EPS)
		out4 = (one_minus_delt*torch.exp(-x**2)) / (erf(y)-erf(x) + EPS)

		# Conditions
		x_cond, y_cond = x.detach(), y.detach() # Make sure that gradients do not go through the conditions

		out1_cond = torch.abs(x_cond - y_cond) &lt; 1e-7 # tensor([False, False, False, False])
		out2_cond = torch.logical_and( torch.logical_and(x_cond&lt;=0, y_cond&lt;=0), torch.logical_not(out1_cond) ) # tensor([False,  True, False, False])
		out3_cond = torch.logical_and( torch.logical_and(x_cond&gt;=0, y_cond&gt;=0), torch.logical_not(out1_cond) ) # tensor([False, False, False,  True]) 
		out4_cond = torch.logical_and( torch.logical_and( torch.logical_not(out1_cond), torch.logical_not(out2_cond) ),
									torch.logical_not(out3_cond) ) # All the other conditions must be false # tensor([ True, False,  True, False])

		# Final expression
		final_out = torch.where(out1_cond,out1,0) + torch.where(out2_cond,out2,0) + torch.where(out3_cond,out3,0) + torch.where(out4_cond,out4,0)

		return final_out
</code></pre>
<p>This works perfectly for the forward pass but, at the backward pass (i.e., for computing the gradients), <strong>I get NaN gradients</strong>. What happens is that some of these expressions result in <em>NaNs</em> or <em>infs</em> for some positions. Nonetheless, the values in these positions are discarded by <em>torch.where</em> (i.e., they are substituted by 0). Despite this, I still get NaN gradients for the final result (<em>final_out</em>), <strong>even though the values which result in NaN gradients are not used in calculating final_out, since torch.where discards them</strong>. In other words, I think that if I were able to <strong>substitute NaN gradients for any other value</strong> (e.g., 0), the gradient calculations would work perfectly.</p>
<p>Any suggestions on how to solve this issue?</p>

username: ptrblck
cooked: <p>Could you post an example input which would yield valid values in the first example but NaNs in the latter?</p>

username: KFrank
cooked: <p>Hi Carlos!</p>
<aside class="quote no-group" data-username="Aeryan" data-post="1" data-topic="185367" 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/aeryan/48/46900_2.png" class="avatar"> Aeryan:</div>
<blockquote>
<p>I still get NaN gradients for the final result (<em>final_out</em>), <strong>even though the values which result in NaN gradients are not used in calculating final_out, since torch.where discards them</strong>.</p>
</blockquote>
</aside>
<p>First, this is a known issue (with no simple fix for <code>torch.where()</code>).  See,<br>
for example, github issues <a href="https://github.com/pytorch/pytorch/issues/68425" rel="noopener nofollow ugc">68425</a> and <a href="https://github.com/pytorch/pytorch/issues/70342" rel="noopener nofollow ugc">70342.</a></p>
<p>Second, I believe that the best fix is to avoid producing <code>nan</code>s (or the values<br>
that subsequently produce <code>nan</code>s) in the forward pass, even in the “branch<br>
not taken” of the <code>torch.where()</code>.  Instead, compute in their places “safe”<br>
(non-<code>nan</code>-producing) values that can be incorrect, because <code>torch.where()</code><br>
<em>will</em> successfully discard the incorrect (but non-<code>nan</code>) gradients that they<br>
lead to.</p>
<p>For a concrete example of how feed “safe” values to the “branch not taken”<br>
in a problem analogous to yours, see this post:</p>
<aside class="quote quote-modified" data-post="4" data-topic="153176">
  <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/k/ecb155/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/unexpected-nan-gradient-propagated-when-applying-boolean-conditions/153176/4">Unexpected NaN gradient propagated when applying boolean conditions</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    Hi rbregier! 

Arul’s explanation is correct.  torch.where() backpropagates a gradient 
of 0 * nan = nan (or maybe 0 * inf = nan) through the “branch not 
taken.” 
An example: 
&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,…
  </blockquote>
</aside>

<p>Best.</p>
<p>K. Frank</p>

username: Aeryan
cooked: <p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a> and <a class="mention" href="/u/kfrank">@KFrank</a> for your answers, and sorry for the late reply.</p>
<p><a class="mention" href="/u/kfrank">@KFrank</a>, I looked into what you mentioned and, yes, that seems to be the root of my problem. In order to solve it, I first tried to use <a href="https://pytorch.org/tutorials/prototype/maskedtensor_overview" rel="noopener nofollow ugc">masked tensors</a> but, unfortunately, they don’t support the <em>erfcx</em> operation (yet). Next, I tried to implement “my own version” of masked tensors by masking the input values for each separate operation so that NaNs are not produced (as you suggested), applying the operations, and then masking the outputs again before adding up the results. This seems to work both for the forward and backward pass. Here is my new implementation of the function <em>_F_1</em>:</p>
<pre><code class="lang-auto">def _F_1(x_, y_):
		# All values in tensor @x must be smaller than values in tensor @y
		x = torch.where(torch.abs(x_)&lt;=torch.abs(y_),x_,y_)
		y = torch.where(torch.abs(y_)&gt;=torch.abs(x_),y_,x_)

		# Obtain masks
		with torch.no_grad():
			out1_cond = torch.abs(x - y) &lt; 1e-7 # tensor([False, False, False, False])
			out2_cond = torch.logical_and( torch.logical_and(x&lt;=0, y&lt;=0), torch.logical_not(out1_cond) ) # tensor([False,  True, False, False])
			out3_cond = torch.logical_and( torch.logical_and(x&gt;=0, y&gt;=0), torch.logical_not(out1_cond) ) # tensor([False, False, False,  True]) 
			out4_cond = torch.logical_and( torch.logical_and( torch.logical_not(out1_cond), torch.logical_not(out2_cond) ),
										   torch.logical_not(out3_cond) ) # All the other conditions must be false


		# Mask input values for each operation
		x1, y1 = where(out1_cond, x, 0), where(out1_cond, y, 0)
		x2, y2 = where(out2_cond, x, -1), where(out2_cond, y, 0)
		x3, y3 = where(out3_cond, x, 0), where(out3_cond, y, 1)
		x4, y4 = where(out4_cond, x, 0), where(out4_cond, y, 1)

		# Apply operations
		delt = ParallelTruncatedGaussian._delta(x, y)
		one_minus_delt = 1 - delt

		out1_m = ParallelTruncatedGaussian._P_1(x1, y1-x1)
		out2_m = one_minus_delt / (delt*erfcx(-y2) - erfcx(-x2))
		out3_m = one_minus_delt / (erfcx(x3) - delt*erfcx(y3))
		out4_m = (one_minus_delt*torch.exp(-x4**2)) / (erf(y4)-erf(x4))

		# Unmask tensors, by setting masked values to 0
		out1 = where(out1_cond, out1_m, 0)
		out2 = where(out2_cond, out2_m, 0)
		out3 = where(out3_cond, out3_m, 0)
		out4 = where(out4_cond, out4_m, 0)

		# Add them up into a single tensor
		final_out = out1 + out2 + out3 + out4

		return final_out
</code></pre>
<p>Also, I have acknowledged your help, <a class="mention" href="/u/kfrank">@KFrank</a> and <a class="mention" href="/u/ptrblck">@ptrblck</a>, in my <a href="https://github.com/TheAeryan/stable-truncated-gaussian/tree/main" rel="noopener nofollow ugc">github repo</a> <img src="https://discuss.pytorch.org/images/emoji/apple/grin.png?v=12" title=":grin:" class="emoji" alt=":grin:" loading="lazy" width="20" height="20">.</p>