'AngleBackward' returned nan values

eagomez: <p>I have a loss function that uses <a href="https://pytorch.org/docs/stable/generated/torch.angle.html" rel="noopener nofollow ugc">torch.angle()</a>, but the argument can take any complex value, therefore even if I do some trick like <code>arg + 1e-7</code> or something along those lines, technically nothing prevents the argument of being exactly the value to produce a <code>nan</code> during the backward pass.</p>
<p>What would be a clean solution to such problem? Can I just plug in <a href="https://pytorch.org/docs/stable/generated/torch.nan_to_num.html" rel="noopener nofollow ugc">torch.nan_to_num()</a> and expect the backward pass to be happy, or should I switch to something like <a href="https://pytorch.org/docs/stable/generated/torch.atan.html#torch.atan" rel="noopener nofollow ugc">torch.atan()</a> and manually add <code>+1e-7</code> to the denominator of the fraction and to replace <a href="https://pytorch.org/docs/stable/generated/torch.angle.html" rel="noopener nofollow ugc">torch.angle()</a> by this? The <code>__call__</code> of this class is the one that is causing me issues:</p>
<p><strong>Edit</strong> I think it is solved now and the issue were denormals that for some reason I forgot they existed. I just set a threshold that corresponds to the same eps that tensorflow uses and seems to be fine, however I found this:<br>
<a href="https://pytorch.org/docs/stable/generated/torch.set_flush_denormal.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.set_flush_denormal.html</a></p>
<p>but it mentions that it is hardware-dependent. What would be a good cross-platform manner to manage this to avoid similar issues? or is it safe to assume that it will be present in most of updated machines?</p>
<pre><code class="lang-auto">class ComplexCompressedMSELoss:
    def __init__(self,
                 c_: float = 0.3,
                 lambda_: float = 0.3,
                 eps: float = 1e-7):
        super().__init__()
        self.c_ = c_
        self.lambda_ = lambda_
        self.eps = eps
 
    def __call__(self, y_pred_mask, x_complex, y_complex):
        # get target magnitude and phase
        y_mag = torch.abs(y_complex)
        y_phase = torch.angle(y_complex)

        # predicted complex stft
        y_pred_mask = y_pred_mask.squeeze(1).permute(0, 2, 1)
        y_pred_complex = y_pred_mask.type(torch.complex64) * x_complex

        # get predicted magnitude and phase
        y_pred_mag = torch.abs(y_pred_complex)
        y_pred_phase = torch.angle(y_pred_complex)

        # target complex exponential        
        y_complex_exp = (y_mag ** self.c_).type(torch.complex64) * \
                torch.exp(1j * y_phase.type(torch.complex64))

        # predicted complex exponential
        y_pred_complex_exp = (y_pred_mag ** self.c_).type(torch.complex64) * \
                torch.exp(1j * y_pred_phase.type(torch.complex64))

        # magnitude only loss component
        mag_loss = torch.abs(y_mag ** self.c_ - y_pred_mag ** self.c_) ** 2
        mag_loss = torch.sum(mag_loss, dim=[1, 2])

        # complex loss component
        complex_loss = torch.abs(y_complex_exp - y_pred_complex_exp) ** 2
        complex_loss = torch.sum(complex_loss, dim=[1, 2])

        # blend both loss components
        loss = (1 - self.lambda_) * mag_loss + (self.lambda_) * complex_loss

        # returns the mean blended loss of the batch
        return torch.mean(loss)
</code></pre>

MinkyuChoi: <p>Hi eagomez,<br>
May I ask how did you solved issue with Nan gradient from <code>AngleBackward</code>?<br>
I am also encountering Nan issue with <code>torch.angle()</code>.<br>
I also tried <code>torch.set_flush_denormal(mode)</code> but the issue still exists.<br>
If you have any other solutions you tried, that would be very helpful for me and others facing the same issue.</p>

eagomez: <p>Hi MinkyuChoi,</p>
<p>I just replaced the potential denormals for a certain threshold that worked in my experiments as follows:</p>
<pre><code class="lang-python">def replace_denormals(x: torch.tensor, threshold=1e-10):
    y = x.clone()
    y[(x &lt; threshold) &amp; (x &gt; -1.0 * threshold)] = threshold
    return y
</code></pre>

MinkyuChoi: <p>Thank you Esteban Gomez, your solution works for me. The Nan issue from the <code>AngleBackward</code> and <code>Torch.Angle()</code> seems gone. However, I still don’t understand why the <code>AngleBackward</code> returns Nan. I will think about it further and will post a new discussion. Thank you again for your help.</p>

eagomez: <p>Glad to know it worked and thanks for the follow-up!</p>

Nitesh_Kumar: <p><code>torch.angle</code> can produce NaN gradients for inputs that are close to (0, 0). <code>torch.atan2</code> produces NaN gradient when the input is exactly (0, 0)<br>
So if you replace <code>torch.angle</code> with <code>torch.atan2</code> then it solves the problem. Both of these do the same thing.<br>
See the example below, I have changed like this and it works<br>
<code>angle  = torch.angle(complx_spect + 1e-7)</code><br>
<code>angle = torch.atan2((complx_spect.imag + 1e-7), (complx_spect.real + 1e-7))</code></p>