Hi Zfzhang!



 zfzhang:

This is also a good solution


Please note – not necessarily relevant to your use case – that
neither your original mask expression nor the more-or-less
equivalent torch.lerp() protects against singularities in the
values (nor gradients).
And torch.where() – perhaps surprisingly – doesn’t protect
against singularities in the gradients.
Consider:
>>> import torch
>>> torch.__version__
'1.9.0'
>>> x = torch.arange (-0.9, 1.0, 0.2, requires_grad = True)
>>> m = (x > 0).float()
>>> y = m * torch.sqrt (x) + (1 - m) * torch.sqrt (-x)
>>> y
tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
       grad_fn=<AddBackward0>)
>>> y.sum().backward()
>>> x.grad
tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan])
>>> x = torch.arange (-0.9, 1.0, 0.2, requires_grad = True)
>>> y = torch.where (x > 0, torch.sqrt (x), torch.sqrt (-x))
>>> y
tensor([0.9486832619, 0.8366600275, 0.7071067691, 0.5477225780, 0.3162277639,
        0.3162277639, 0.5477225780, 0.7071067691, 0.8366600275, 0.9486832619],
       grad_fn=<SWhereBackward>)
>>> y.sum().backward()
>>> x.grad
tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan])
>>> x = torch.arange (-0.9, 1.0, 0.2, requires_grad = True)
>>> y = torch.sqrt (x.abs())
>>> y
tensor([0.9486832619, 0.8366600275, 0.7071067691, 0.5477225780, 0.3162277639,
        0.3162277639, 0.5477225780, 0.7071067691, 0.8366600275, 0.9486832619],
       grad_fn=<SqrtBackward>)
>>> y.sum().backward()
>>> x.grad
tensor([-0.5270463228, -0.5976142883, -0.7071067691, -0.9128708839,
        -1.5811388493,  1.5811388493,  0.9128708839,  0.7071067691,
         0.5976142883,  0.5270463228])

Best.
K. Frank