torch.norm is numerically unstable at zero for multidim reductions #37323
Description
Note the inconsistent gradient computation:
import torch
a = torch.zeros(3, 3, 3, requires_grad = True)
print(torch.autograd.grad(a.norm(dim = 1).sum(), (a,))[0])
print(torch.autograd.grad(a.norm(dim = (1, 2)).sum(), (a,))[0])
print(torch.__version__)
tensor([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]])
tensor([[[nan, nan, nan],
[nan, nan, nan],
[nan, nan, nan]],
[[nan, nan, nan],
[nan, nan, nan],
[nan, nan, nan]],
[[nan, nan, nan],
[nan, nan, nan],
[nan, nan, nan]]])
1.6.0.dev20200417
Relevant code: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/LinearAlgebra.cpp#L558
Relevant discussion: #37272 (comment)