Incorrect gradient calculation with torch.where and NaNs

[KFrank]

Hi Carlos!

@Aeryan

I still get NaN gradients for the final result (final_out), even though the values which result in NaN gradients are not used in calculating final_out, since torch.where discards them.

First, this is a known issue (with no simple fix for torch.where()). See,
for example, github issues 68425 and 70342.

Second, I believe that the best fix is to avoid producing nans (or the values
that subsequently produce nans) in the forward pass, even in the “branch
not taken” of the torch.where(). Instead, compute in their places “safe”
(non-nan-producing) values that can be incorrect, because torch.where()
will successfully discard the incorrect (but non-nan) gradients that they
lead to.

For a concrete example of how feed “safe” values to the “branch not taken”
in a problem analogous to yours, see this post:

Unexpected NaN gradient propagated when applying boolean conditions

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: >>> import torch >>> torch.__version__ '1.10.2' >>> s1 = torch.zeros (4, requires_grad = True) >>> s2 = torch.zeros (4, requires_grad = True) >>> t1 = torch.tensor ([0.5, 0.5, 0.0, 0.0]) >>> t2 = torch.tensor ([0.5, 0.5, 1.e-7, 1.e-7]) >>> torch.where (t1 > 1.e-6, s1 / t1, s1).sum().backward() >>> torch.where (t2 > 1.e-6,…

Best.

K. Frank