Repository: pytorch/pytorch
Pull request: Norm subgradient at 0
Fix #2421 

Note that with this, at `0`, it is now different to use `.norm()` or `.pow(2).sum().sqrt()` as the first one will return `0` and the second one `NaN`

Problem statement:
Gradient of zero norm is nan
If a norm is zero, its gradient returns nan:

```python
x = Variable(torch.zeros(1), requires_grad=True)
x.norm().backward()
print x.grad

# Variable containing:
# nan
# [torch.FloatTensor of size 1]
```

Obviously just happening because the gradient divides by the norm, but the (sub)gradient here should probably be zero, or at least not nan, since that will propagate to make all updates nan.  Probably low priority, as it's not going to be an issue in 99% of cases, but we're doing a few things with (exact) line searches where this caused a nan to appear, breaking everything downstream.

Hints from issue discussion:
I'm encountering exactly the same issue! Spent hours on debugging, just to find PyTorch has a bug on such basic thing.
+1 just found this bug too
+1 for this bug. Temporarily changing my code to something like the following for the sake of debugging.

    x = Variable(torch.zeros(1), requires_grad=True)
    y = x + 1e-16
    y.norm().backward()
    print x.grad

The thing is that in the 2 norm, there is a square root, which has a gradient of `+ infinity` at 0.
The gradient gives you `nan` because you then multiply 0 and an infinity during the backward pass.
For a scalar, norm 2 is basically abs. But x.abs().backward() gives you 0 gradient. In this sense, it's not coherent.
I found this error, too

Commit message: add test for subgradient

Patch to review:
diff --git a/test/test_autograd.py b/test/test_autograd.py
index e90334eb8816..e2ba2c66b061 100644
--- a/test/test_autograd.py
+++ b/test/test_autograd.py
@@ -1565,6 +1565,17 @@ def get_out_and_unrefed_cycle():
         # This will segfault if things have been erroneously released
         out.backward(torch.randn(out.size()))
 
+    def test_norm_subgradient(self):
+        def run_test(input_size, norm_deg):
+            input = Variable(torch.randn(*input_size).clamp(min=0), requires_grad=True)
+            out = input.norm(norm_deg)
+            out.backward()
+            self.assertEqual(input.grad.data[input.data == 0].abs().sum(), 0)
+
+        run_test((10,), 2)
+        run_test((10, 10), 2)
+        run_test((10,), 3)
+
 
 def index_variable(shape, max_indices):
     if not isinstance(shape, tuple):
diff --git a/torch/autograd/_functions/reduce.py b/torch/autograd/_functions/reduce.py
index 0a6f37af8f0b..199b87c8b931 100644
--- a/torch/autograd/_functions/reduce.py
+++ b/torch/autograd/_functions/reduce.py
@@ -226,43 +226,35 @@ def forward(ctx, input, p=2, dim=None, keepdim=None):
         ctx.keepdim = False if keepdim is None else keepdim
 
         if dim is None:
-            ctx.norm = input.norm(p)
-            ctx.save_for_backward(input)
-            return input.new((ctx.norm,))
+            norm = input.norm(p)
+            output = input.new((norm,))
         else:
             if keepdim is not None:
                 output = input.norm(p, dim, keepdim=keepdim)
             else:
                 output = input.norm(p, dim)
-            ctx.save_for_backward(input, output)
-            return output
+        ctx.save_for_backward(input, output)
+        return output
 
     @staticmethod
     def backward(ctx, grad_output):
-        if ctx.dim is None:
-            input, = ctx.saved_variables
-            if ctx.p == 2:
-                scale_v = (grad_output / ctx.norm).expand_as(input)
-                return input.mul(scale_v), None, None, None
-            else:
-                pow = input.abs().pow(ctx.p - 2)
-                scale_v = (grad_output / ctx.norm ** (ctx.p - 1)).expand_as(input)
-                return input.mul(pow).mul(scale_v), None, None, None
+        input, output = ctx.saved_variables
+        if ctx.dim is not None and ctx.keepdim is False and input.dim() != 1:
+            grad_output = grad_output.unsqueeze(ctx.dim)
+            output = output.unsqueeze(ctx.dim)
+
+        if ctx.p == 2:
+            grad_input = input.mul(grad_output).div(output)
         else:
-            input, output = ctx.saved_variables
+            input_pow = input.abs().pow(ctx.p - 2)
+            output_pow = output.pow(ctx.p - 1)
+            grad_input = input.mul(input_pow).mul(grad_output).div(output_pow)
 
-            if ctx.keepdim is False and input.dim() != 1:
-                grad_output = grad_output.unsqueeze(ctx.dim)
-                output = output.unsqueeze(ctx.dim)
+        # Special case at 0 where we return a subgradient containing 0
+        zero_mask = (output == 0).type_as(grad_input)
+        grad_input.mul(1 - zero_mask)
 
-            big_grad_output = grad_output.expand_as(input)
-            if ctx.p == 2:
-                big_output = output.expand_as(input)
-                return input.mul(big_grad_output).div(big_output), None, None, None
-            else:
-                pow = input.abs().pow(ctx.p - 2)
-                big_output = output.pow(ctx.p - 1).expand_as(input)
-                return input.mul(pow).mul(big_grad_output).div(big_output), None, None, None
+        return grad_input, None, None, None
 
 
 # TODO: renorm


Code review comments:
diff_hunk: @@ -226,43 +226,35 @@ def forward(ctx, input, p=2, dim=None, keepdim=None):
         ctx.keepdim = False if keepdim is None else keepdim
 
         if dim is None:
-            ctx.norm = input.norm(p)
-            ctx.save_for_backward(input)
-            return input.new((ctx.norm,))
+            norm = input.norm(p)
+            output = input.new((norm,))
         else:
             if keepdim is not None:
                 output = input.norm(p, dim, keepdim=keepdim)
             else:
                 output = input.norm(p, dim)
-            ctx.save_for_backward(input, output)
-            return output
+        ctx.save_for_backward(input, output)
+        return output
 
     @staticmethod
     def backward(ctx, grad_output):
-        if ctx.dim is None:
-            input, = ctx.saved_variables
-            if ctx.p == 2:
-                scale_v = (grad_output / ctx.norm).expand_as(input)
-                return input.mul(scale_v), None, None, None
-            else:
-                pow = input.abs().pow(ctx.p - 2)
-                scale_v = (grad_output / ctx.norm ** (ctx.p - 1)).expand_as(input)
-                return input.mul(pow).mul(scale_v), None, None, None
+        input, output = ctx.saved_variables
+        if ctx.dim is not None and ctx.keepdim is False and input.dim() != 1:
+            grad_output = grad_output.unsqueeze(ctx.dim)
+            output = output.unsqueeze(ctx.dim)
+
+        if ctx.p == 2:
+            grad_input = input.mul(grad_output).div(output)
         else:
-            input, output = ctx.saved_variables
+            input_pow = input.abs().pow(ctx.p - 2)
+            output_pow = output.pow(ctx.p - 1)
+            grad_input = input.mul(input_pow).mul(grad_output).div(output_pow)
 
-            if ctx.keepdim is False and input.dim() != 1:
-                grad_output = grad_output.unsqueeze(ctx.dim)
-                output = output.unsqueeze(ctx.dim)
+        # Special case at 0 where we return a subgradient containing 0
+        zero_mask = (output == 0).type_as(grad_input)
+        grad_input.mul(1 - zero_mask)
path: torch/autograd/_functions/reduce.py
text: @user1:
Did you intent to do inplace multiplication?

@author:
Yes, my bad, the test was wrong... Thanks for spotting that.