We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
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.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torch/autograd/_functions/reduce.py`
Content:
```
1 from functools import reduce
2 
3 from ..function import Function
4 from ..variable import Variable
5 import torch
6 
7 
8 class Sum(Function):
9 
10     @staticmethod
11     def forward(ctx, input, dim=None, keepdim=None):
12         ctx.dim = dim
13         ctx.keepdim = False if keepdim is None else keepdim
14         ctx.input_size = input.size()
15         if dim is None:
16             return input.new((input.sum(),))
17         else:
18             if keepdim is not None:
19                 return input.sum(dim, keepdim=keepdim)
20             else:
21                 return input.sum(dim)
22 
23     @staticmethod
24     def backward(ctx, grad_output):
25         if ctx.dim is None:
26             return grad_output.expand(ctx.input_size), None, None
27         else:
28             if ctx.keepdim is False and len(ctx.input_size) != 1:
29                 grad_output = grad_output.unsqueeze(ctx.dim)
30 
31             repeats = [1 for _ in ctx.input_size]
32             repeats[ctx.dim] = ctx.input_size[ctx.dim]
33             return grad_output.repeat(*repeats), None, None
34 
35 
36 class Prod(Function):
37 
38     @staticmethod
39     def forward(ctx, input, dim=None, keepdim=None):
40         ctx.dim = dim
41         ctx.keepdim = False if keepdim is None else keepdim
42         ctx.input_size = input.size()
43         if dim is None:
44             ctx.result = input.prod()
45             ctx.save_for_backward(input)
46             return input.new((ctx.result,))
47         else:
48             if keepdim is not None:
49                 output = input.prod(dim, keepdim=keepdim)
50             else:
51                 output = input.prod(dim)
52             ctx.save_for_backward(input, output)
53             return output
54 
55     @staticmethod
56     def backward(ctx, grad_output):
57         def safe_zeros_backward(inp, dim):
58             # note that the gradient is equivalent to:
59             # cumprod(exclusive, normal) * cumprod(exclusive, reverse), e.g.:
60             # input:                        [    a,     b,     c]
61             # cumprod(exclusive, normal):   [1    ,     a, a * b]
62             # cumprod(exclusive, reverse):  [b * c,     c,     1]
63             # product:                      [b * c, a * c, a * b]
64             # and this is safe under input with 0s.
65             if inp.size(dim) == 1:
66                 return grad_output
67 
68             ones_size = torch.Size((inp.size()[:dim] + (1,) + inp.size()[dim + 1:]))
69             ones = Variable(grad_output.data.new(ones_size).fill_(1))
70             exclusive_normal_nocp = torch.cat((ones, inp.narrow(dim, 0, inp.size(dim) - 1)), dim)
71             exclusive_normal = exclusive_normal_nocp.cumprod(dim)
72 
73             def reverse_dim(var, dim):
74                 index = Variable(torch.arange(var.size(dim) - 1, -1, -1, out=var.data.new().long()))
75                 return var.index_select(dim, index)
76 
77             narrow_reverse = reverse_dim(inp.narrow(dim, 1, inp.size(dim) - 1), dim)
78             exclusive_reverse_nocp = torch.cat((ones, narrow_reverse), dim)
79             exclusive_reverse = reverse_dim(exclusive_reverse_nocp.cumprod(dim), dim)
80 
81             grad_input = grad_output.expand_as(exclusive_normal).mul(exclusive_normal.mul(exclusive_reverse))
82             return grad_input
83 
84         if ctx.dim is None:
85             input, = ctx.saved_variables
86             zero_idx = (input.data == 0).nonzero()
87             if zero_idx.dim() == 0:
88                 return grad_output.mul(ctx.result).expand_as(input).div(input), None, None
89             elif zero_idx.size(0) > 1:
90                 return (grad_output * 0).expand_as(input), None, None
91             else:
92                 return safe_zeros_backward(input.contiguous().view(-1), 0).view_as(input), None, None
93 
94         else:
95             input, output = ctx.saved_variables
96             dim = ctx.dim if ctx.dim >= 0 else ctx.dim + input.dim()
97             if ctx.keepdim is False and len(ctx.input_size) != 1:
98                 grad_output = grad_output.unsqueeze(dim)
99                 output = output.unsqueeze(dim)
100 
101             zero_mask = input == 0
102             slice_zero_count = zero_mask.sum(dim, True)
103             total_zeros = slice_zero_count.data.sum()
104             if total_zeros == 0:
105                 grad_input = grad_output.mul(output).expand_as(input).div(input)
106             else:
107                 grad_input = safe_zeros_backward(input, dim)
108 
109             return grad_input, None, None
110 
111 
112 class Mean(Function):
113 
114     @staticmethod
115     def forward(ctx, input, dim=None, keepdim=None):
116         ctx.dim = dim
117         ctx.keepdim = False if keepdim is None else keepdim
118         ctx.input_size = input.size()
119         if dim is None:
120             return input.new((input.mean(),))
121         else:
122             if keepdim is not None:
123                 return input.mean(dim, keepdim=keepdim)
124             else:
125                 return input.mean(dim)
126 
127     @staticmethod
128     def backward(ctx, grad_output):
129         if ctx.dim is None:
130             grad_input_val = grad_output / reduce(lambda x, y: x * y, ctx.input_size, 1)
131             return grad_input_val.expand(ctx.input_size), None, None
132         else:
133             if ctx.keepdim is False and len(ctx.input_size) != 1:
134                 grad_output = grad_output.unsqueeze(ctx.dim)
135 
136             repeats = [1 for _ in ctx.input_size]
137             dim_size = ctx.input_size[ctx.dim]
138             repeats[ctx.dim] = dim_size
139             return grad_output.repeat(*repeats).div_(dim_size), None, None
140 
141 
142 class _SelectionFunction(Function):
143     has_all_reduce = True
144     # additional_args is prepended before dim when calling the tensor
145     # function. It's a no-op for subclasses other than kthvalue.
146     # kthvalue not only requires us to pass a dim, but also precede it with k.
147 
148     @classmethod
149     def forward(cls, ctx, input, dim=None, keepdim=None, additional_args=tuple()):
150         fn = getattr(input, cls.__name__.lower())
151         ctx.dim = dim
152         ctx.keepdim = False if keepdim is None else keepdim
153         ctx.additional_args = additional_args
154         ctx.input_size = input.size()
155         if ctx.dim is None and cls.has_all_reduce:
156             value = fn(*additional_args)
157             ctx.indices_tuple = tuple(input.eq(value).nonzero()[0])
158             return input.new((value,))
159         else:
160             if ctx.dim is None:
161                 dim = input.dim() - 1
162             else:
163                 dim = ctx.dim
164             args = (dim,)
165             if additional_args:
166                 args = additional_args + args
167             if keepdim is not None:
168                 output, indices = fn(*args, keepdim=keepdim)
169             else:
170                 output, indices = fn(*args)
171             ctx.save_for_backward(indices)
172             ctx.mark_non_differentiable(indices)
173             return output, indices
174 
175     @classmethod
176     def backward(cls, ctx, grad_output, grad_indices=None):
177         grad_input = Variable(grad_output.data.new(*ctx.input_size).zero_())
178         if ctx.dim is None and cls.has_all_reduce:
179             grad_input[ctx.indices_tuple] = grad_output
180         else:
181             if ctx.dim is None:
182                 dim = len(ctx.input_size) - 1
183             else:
184                 dim = ctx.dim
185 
186             indices, = ctx.saved_variables
187             if ctx.keepdim is False and len(ctx.input_size) != 1:
188                 grad_output = grad_output.unsqueeze(dim)
189                 grad_indices = grad_indices.unsqueeze(dim)
190                 indices = indices.unsqueeze(dim)
191 
192             grad_input.scatter_(dim, indices, grad_output)
193         return grad_input, None, None, None
194 
195 
196 class Max(_SelectionFunction):
197     pass
198 
199 
200 class Min(_SelectionFunction):
201     pass
202 
203 
204 class Mode(_SelectionFunction):
205     has_all_reduce = False
206 
207 
208 class Median(_SelectionFunction):
209     pass
210 
211 
212 class Kthvalue(_SelectionFunction):
213     has_all_reduce = False
214 
215     @classmethod
216     def forward(cls, ctx, input, k, dim=None, keepdim=None):
217         return super(Kthvalue, cls).forward(ctx, input, dim, keepdim, (k,))
218 
219 
220 class Norm(Function):
221 
222     @staticmethod
223     def forward(ctx, input, p=2, dim=None, keepdim=None):
224         ctx.p = p
225         ctx.dim = dim
226         ctx.keepdim = False if keepdim is None else keepdim
227 
228         if dim is None:
229             ctx.norm = input.norm(p)
230             ctx.save_for_backward(input)
231             return input.new((ctx.norm,))
232         else:
233             if keepdim is not None:
234                 output = input.norm(p, dim, keepdim=keepdim)
235             else:
236                 output = input.norm(p, dim)
237             ctx.save_for_backward(input, output)
238             return output
239 
240     @staticmethod
241     def backward(ctx, grad_output):
242         if ctx.dim is None:
243             input, = ctx.saved_variables
244             if ctx.p == 2:
245                 scale_v = (grad_output / ctx.norm).expand_as(input)
246                 return input.mul(scale_v), None, None, None
247             else:
248                 pow = input.abs().pow(ctx.p - 2)
249                 scale_v = (grad_output / ctx.norm ** (ctx.p - 1)).expand_as(input)
250                 return input.mul(pow).mul(scale_v), None, None, None
251         else:
252             input, output = ctx.saved_variables
253 
254             if ctx.keepdim is False and input.dim() != 1:
255                 grad_output = grad_output.unsqueeze(ctx.dim)
256                 output = output.unsqueeze(ctx.dim)
257 
258             big_grad_output = grad_output.expand_as(input)
259             if ctx.p == 2:
260                 big_output = output.expand_as(input)
261                 return input.mul(big_grad_output).div(big_output), None, None, None
262             else:
263                 pow = input.abs().pow(ctx.p - 2)
264                 big_output = output.pow(ctx.p - 1).expand_as(input)
265                 return input.mul(pow).mul(big_grad_output).div(big_output), None, None, None
266 
267 
268 # TODO: renorm
269 # TODO: std
270 # TODO: var
271
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torch/autograd/_functions/reduce.py b/torch/autograd/_functions/reduce.py
--- a/torch/autograd/_functions/reduce.py
+++ b/torch/autograd/_functions/reduce.py
@@ -226,43 +226,34 @@
         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
+        grad_input.masked_fill_(output == 0, 0)
 
-            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

{"golden_diff": "diff --git a/torch/autograd/_functions/reduce.py b/torch/autograd/_functions/reduce.py\n--- a/torch/autograd/_functions/reduce.py\n+++ b/torch/autograd/_functions/reduce.py\n@@ -226,43 +226,34 @@\n         ctx.keepdim = False if keepdim is None else keepdim\n \n         if dim is None:\n-            ctx.norm = input.norm(p)\n-            ctx.save_for_backward(input)\n-            return input.new((ctx.norm,))\n+            norm = input.norm(p)\n+            output = input.new((norm,))\n         else:\n             if keepdim is not None:\n                 output = input.norm(p, dim, keepdim=keepdim)\n             else:\n                 output = input.norm(p, dim)\n-            ctx.save_for_backward(input, output)\n-            return output\n+        ctx.save_for_backward(input, output)\n+        return output\n \n     @staticmethod\n     def backward(ctx, grad_output):\n-        if ctx.dim is None:\n-            input, = ctx.saved_variables\n-            if ctx.p == 2:\n-                scale_v = (grad_output / ctx.norm).expand_as(input)\n-                return input.mul(scale_v), None, None, None\n-            else:\n-                pow = input.abs().pow(ctx.p - 2)\n-                scale_v = (grad_output / ctx.norm ** (ctx.p - 1)).expand_as(input)\n-                return input.mul(pow).mul(scale_v), None, None, None\n+        input, output = ctx.saved_variables\n+        if ctx.dim is not None and ctx.keepdim is False and input.dim() != 1:\n+            grad_output = grad_output.unsqueeze(ctx.dim)\n+            output = output.unsqueeze(ctx.dim)\n+\n+        if ctx.p == 2:\n+            grad_input = input.mul(grad_output).div(output)\n         else:\n-            input, output = ctx.saved_variables\n+            input_pow = input.abs().pow(ctx.p - 2)\n+            output_pow = output.pow(ctx.p - 1)\n+            grad_input = input.mul(input_pow).mul(grad_output).div(output_pow)\n \n-            if ctx.keepdim is False and input.dim() != 1:\n-                grad_output = grad_output.unsqueeze(ctx.dim)\n-                output = output.unsqueeze(ctx.dim)\n+        # Special case at 0 where we return a subgradient containing 0\n+        grad_input.masked_fill_(output == 0, 0)\n \n-            big_grad_output = grad_output.expand_as(input)\n-            if ctx.p == 2:\n-                big_output = output.expand_as(input)\n-                return input.mul(big_grad_output).div(big_output), None, None, None\n-            else:\n-                pow = input.abs().pow(ctx.p - 2)\n-                big_output = output.pow(ctx.p - 1).expand_as(input)\n-                return input.mul(pow).mul(big_grad_output).div(big_output), None, None, None\n+        return grad_input, None, None, None\n \n \n # TODO: renorm\n", "issue": "Gradient of zero norm is nan\nIf a norm is zero, its gradient returns nan:\r\n\r\n```python\r\nx = Variable(torch.zeros(1), requires_grad=True)\r\nx.norm().backward()\r\nprint x.grad\r\n\r\n# Variable containing:\r\n# nan\r\n# [torch.FloatTensor of size 1]\r\n```\r\n\r\nObviously 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.\r\n\n", "before_files": [{"content": "from functools import reduce\n\nfrom ..function import Function\nfrom ..variable import Variable\nimport torch\n\n\nclass Sum(Function):\n\n    @staticmethod\n    def forward(ctx, input, dim=None, keepdim=None):\n        ctx.dim = dim\n        ctx.keepdim = False if keepdim is None else keepdim\n        ctx.input_size = input.size()\n        if dim is None:\n            return input.new((input.sum(),))\n        else:\n            if keepdim is not None:\n                return input.sum(dim, keepdim=keepdim)\n            else:\n                return input.sum(dim)\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        if ctx.dim is None:\n            return grad_output.expand(ctx.input_size), None, None\n        else:\n            if ctx.keepdim is False and len(ctx.input_size) != 1:\n                grad_output = grad_output.unsqueeze(ctx.dim)\n\n            repeats = [1 for _ in ctx.input_size]\n            repeats[ctx.dim] = ctx.input_size[ctx.dim]\n            return grad_output.repeat(*repeats), None, None\n\n\nclass Prod(Function):\n\n    @staticmethod\n    def forward(ctx, input, dim=None, keepdim=None):\n        ctx.dim = dim\n        ctx.keepdim = False if keepdim is None else keepdim\n        ctx.input_size = input.size()\n        if dim is None:\n            ctx.result = input.prod()\n            ctx.save_for_backward(input)\n            return input.new((ctx.result,))\n        else:\n            if keepdim is not None:\n                output = input.prod(dim, keepdim=keepdim)\n            else:\n                output = input.prod(dim)\n            ctx.save_for_backward(input, output)\n            return output\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        def safe_zeros_backward(inp, dim):\n            # note that the gradient is equivalent to:\n            # cumprod(exclusive, normal) * cumprod(exclusive, reverse), e.g.:\n            # input:                        [    a,     b,     c]\n            # cumprod(exclusive, normal):   [1    ,     a, a * b]\n            # cumprod(exclusive, reverse):  [b * c,     c,     1]\n            # product:                      [b * c, a * c, a * b]\n            # and this is safe under input with 0s.\n            if inp.size(dim) == 1:\n                return grad_output\n\n            ones_size = torch.Size((inp.size()[:dim] + (1,) + inp.size()[dim + 1:]))\n            ones = Variable(grad_output.data.new(ones_size).fill_(1))\n            exclusive_normal_nocp = torch.cat((ones, inp.narrow(dim, 0, inp.size(dim) - 1)), dim)\n            exclusive_normal = exclusive_normal_nocp.cumprod(dim)\n\n            def reverse_dim(var, dim):\n                index = Variable(torch.arange(var.size(dim) - 1, -1, -1, out=var.data.new().long()))\n                return var.index_select(dim, index)\n\n            narrow_reverse = reverse_dim(inp.narrow(dim, 1, inp.size(dim) - 1), dim)\n            exclusive_reverse_nocp = torch.cat((ones, narrow_reverse), dim)\n            exclusive_reverse = reverse_dim(exclusive_reverse_nocp.cumprod(dim), dim)\n\n            grad_input = grad_output.expand_as(exclusive_normal).mul(exclusive_normal.mul(exclusive_reverse))\n            return grad_input\n\n        if ctx.dim is None:\n            input, = ctx.saved_variables\n            zero_idx = (input.data == 0).nonzero()\n            if zero_idx.dim() == 0:\n                return grad_output.mul(ctx.result).expand_as(input).div(input), None, None\n            elif zero_idx.size(0) > 1:\n                return (grad_output * 0).expand_as(input), None, None\n            else:\n                return safe_zeros_backward(input.contiguous().view(-1), 0).view_as(input), None, None\n\n        else:\n            input, output = ctx.saved_variables\n            dim = ctx.dim if ctx.dim >= 0 else ctx.dim + input.dim()\n            if ctx.keepdim is False and len(ctx.input_size) != 1:\n                grad_output = grad_output.unsqueeze(dim)\n                output = output.unsqueeze(dim)\n\n            zero_mask = input == 0\n            slice_zero_count = zero_mask.sum(dim, True)\n            total_zeros = slice_zero_count.data.sum()\n            if total_zeros == 0:\n                grad_input = grad_output.mul(output).expand_as(input).div(input)\n            else:\n                grad_input = safe_zeros_backward(input, dim)\n\n            return grad_input, None, None\n\n\nclass Mean(Function):\n\n    @staticmethod\n    def forward(ctx, input, dim=None, keepdim=None):\n        ctx.dim = dim\n        ctx.keepdim = False if keepdim is None else keepdim\n        ctx.input_size = input.size()\n        if dim is None:\n            return input.new((input.mean(),))\n        else:\n            if keepdim is not None:\n                return input.mean(dim, keepdim=keepdim)\n            else:\n                return input.mean(dim)\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        if ctx.dim is None:\n            grad_input_val = grad_output / reduce(lambda x, y: x * y, ctx.input_size, 1)\n            return grad_input_val.expand(ctx.input_size), None, None\n        else:\n            if ctx.keepdim is False and len(ctx.input_size) != 1:\n                grad_output = grad_output.unsqueeze(ctx.dim)\n\n            repeats = [1 for _ in ctx.input_size]\n            dim_size = ctx.input_size[ctx.dim]\n            repeats[ctx.dim] = dim_size\n            return grad_output.repeat(*repeats).div_(dim_size), None, None\n\n\nclass _SelectionFunction(Function):\n    has_all_reduce = True\n    # additional_args is prepended before dim when calling the tensor\n    # function. It's a no-op for subclasses other than kthvalue.\n    # kthvalue not only requires us to pass a dim, but also precede it with k.\n\n    @classmethod\n    def forward(cls, ctx, input, dim=None, keepdim=None, additional_args=tuple()):\n        fn = getattr(input, cls.__name__.lower())\n        ctx.dim = dim\n        ctx.keepdim = False if keepdim is None else keepdim\n        ctx.additional_args = additional_args\n        ctx.input_size = input.size()\n        if ctx.dim is None and cls.has_all_reduce:\n            value = fn(*additional_args)\n            ctx.indices_tuple = tuple(input.eq(value).nonzero()[0])\n            return input.new((value,))\n        else:\n            if ctx.dim is None:\n                dim = input.dim() - 1\n            else:\n                dim = ctx.dim\n            args = (dim,)\n            if additional_args:\n                args = additional_args + args\n            if keepdim is not None:\n                output, indices = fn(*args, keepdim=keepdim)\n            else:\n                output, indices = fn(*args)\n            ctx.save_for_backward(indices)\n            ctx.mark_non_differentiable(indices)\n            return output, indices\n\n    @classmethod\n    def backward(cls, ctx, grad_output, grad_indices=None):\n        grad_input = Variable(grad_output.data.new(*ctx.input_size).zero_())\n        if ctx.dim is None and cls.has_all_reduce:\n            grad_input[ctx.indices_tuple] = grad_output\n        else:\n            if ctx.dim is None:\n                dim = len(ctx.input_size) - 1\n            else:\n                dim = ctx.dim\n\n            indices, = ctx.saved_variables\n            if ctx.keepdim is False and len(ctx.input_size) != 1:\n                grad_output = grad_output.unsqueeze(dim)\n                grad_indices = grad_indices.unsqueeze(dim)\n                indices = indices.unsqueeze(dim)\n\n            grad_input.scatter_(dim, indices, grad_output)\n        return grad_input, None, None, None\n\n\nclass Max(_SelectionFunction):\n    pass\n\n\nclass Min(_SelectionFunction):\n    pass\n\n\nclass Mode(_SelectionFunction):\n    has_all_reduce = False\n\n\nclass Median(_SelectionFunction):\n    pass\n\n\nclass Kthvalue(_SelectionFunction):\n    has_all_reduce = False\n\n    @classmethod\n    def forward(cls, ctx, input, k, dim=None, keepdim=None):\n        return super(Kthvalue, cls).forward(ctx, input, dim, keepdim, (k,))\n\n\nclass Norm(Function):\n\n    @staticmethod\n    def forward(ctx, input, p=2, dim=None, keepdim=None):\n        ctx.p = p\n        ctx.dim = dim\n        ctx.keepdim = False if keepdim is None else keepdim\n\n        if dim is None:\n            ctx.norm = input.norm(p)\n            ctx.save_for_backward(input)\n            return input.new((ctx.norm,))\n        else:\n            if keepdim is not None:\n                output = input.norm(p, dim, keepdim=keepdim)\n            else:\n                output = input.norm(p, dim)\n            ctx.save_for_backward(input, output)\n            return output\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        if ctx.dim is None:\n            input, = ctx.saved_variables\n            if ctx.p == 2:\n                scale_v = (grad_output / ctx.norm).expand_as(input)\n                return input.mul(scale_v), None, None, None\n            else:\n                pow = input.abs().pow(ctx.p - 2)\n                scale_v = (grad_output / ctx.norm ** (ctx.p - 1)).expand_as(input)\n                return input.mul(pow).mul(scale_v), None, None, None\n        else:\n            input, output = ctx.saved_variables\n\n            if ctx.keepdim is False and input.dim() != 1:\n                grad_output = grad_output.unsqueeze(ctx.dim)\n                output = output.unsqueeze(ctx.dim)\n\n            big_grad_output = grad_output.expand_as(input)\n            if ctx.p == 2:\n                big_output = output.expand_as(input)\n                return input.mul(big_grad_output).div(big_output), None, None, None\n            else:\n                pow = input.abs().pow(ctx.p - 2)\n                big_output = output.pow(ctx.p - 1).expand_as(input)\n                return input.mul(pow).mul(big_grad_output).div(big_output), None, None, None\n\n\n# TODO: renorm\n# TODO: std\n# TODO: var\n", "path": "torch/autograd/_functions/reduce.py"}], "after_files": [{"content": "from functools import reduce\n\nfrom ..function import Function\nfrom ..variable import Variable\nimport torch\n\n\nclass Sum(Function):\n\n    @staticmethod\n    def forward(ctx, input, dim=None, keepdim=None):\n        ctx.dim = dim\n        ctx.keepdim = False if keepdim is None else keepdim\n        ctx.input_size = input.size()\n        if dim is None:\n            return input.new((input.sum(),))\n        else:\n            if keepdim is not None:\n                return input.sum(dim, keepdim=keepdim)\n            else:\n                return input.sum(dim)\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        if ctx.dim is None:\n            return grad_output.expand(ctx.input_size), None, None\n        else:\n            if ctx.keepdim is False and len(ctx.input_size) != 1:\n                grad_output = grad_output.unsqueeze(ctx.dim)\n\n            repeats = [1 for _ in ctx.input_size]\n            repeats[ctx.dim] = ctx.input_size[ctx.dim]\n            return grad_output.repeat(*repeats), None, None\n\n\nclass Prod(Function):\n\n    @staticmethod\n    def forward(ctx, input, dim=None, keepdim=None):\n        ctx.dim = dim\n        ctx.keepdim = False if keepdim is None else keepdim\n        ctx.input_size = input.size()\n        if dim is None:\n            ctx.result = input.prod()\n            ctx.save_for_backward(input)\n            return input.new((ctx.result,))\n        else:\n            if keepdim is not None:\n                output = input.prod(dim, keepdim=keepdim)\n            else:\n                output = input.prod(dim)\n            ctx.save_for_backward(input, output)\n            return output\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        def safe_zeros_backward(inp, dim):\n            # note that the gradient is equivalent to:\n            # cumprod(exclusive, normal) * cumprod(exclusive, reverse), e.g.:\n            # input:                        [    a,     b,     c]\n            # cumprod(exclusive, normal):   [1    ,     a, a * b]\n            # cumprod(exclusive, reverse):  [b * c,     c,     1]\n            # product:                      [b * c, a * c, a * b]\n            # and this is safe under input with 0s.\n            if inp.size(dim) == 1:\n                return grad_output\n\n            ones_size = torch.Size((inp.size()[:dim] + (1,) + inp.size()[dim + 1:]))\n            ones = Variable(grad_output.data.new(ones_size).fill_(1))\n            exclusive_normal_nocp = torch.cat((ones, inp.narrow(dim, 0, inp.size(dim) - 1)), dim)\n            exclusive_normal = exclusive_normal_nocp.cumprod(dim)\n\n            def reverse_dim(var, dim):\n                index = Variable(torch.arange(var.size(dim) - 1, -1, -1, out=var.data.new().long()))\n                return var.index_select(dim, index)\n\n            narrow_reverse = reverse_dim(inp.narrow(dim, 1, inp.size(dim) - 1), dim)\n            exclusive_reverse_nocp = torch.cat((ones, narrow_reverse), dim)\n            exclusive_reverse = reverse_dim(exclusive_reverse_nocp.cumprod(dim), dim)\n\n            grad_input = grad_output.expand_as(exclusive_normal).mul(exclusive_normal.mul(exclusive_reverse))\n            return grad_input\n\n        if ctx.dim is None:\n            input, = ctx.saved_variables\n            zero_idx = (input.data == 0).nonzero()\n            if zero_idx.dim() == 0:\n                return grad_output.mul(ctx.result).expand_as(input).div(input), None, None\n            elif zero_idx.size(0) > 1:\n                return (grad_output * 0).expand_as(input), None, None\n            else:\n                return safe_zeros_backward(input.contiguous().view(-1), 0).view_as(input), None, None\n\n        else:\n            input, output = ctx.saved_variables\n            dim = ctx.dim if ctx.dim >= 0 else ctx.dim + input.dim()\n            if ctx.keepdim is False and len(ctx.input_size) != 1:\n                grad_output = grad_output.unsqueeze(dim)\n                output = output.unsqueeze(dim)\n\n            zero_mask = input == 0\n            slice_zero_count = zero_mask.sum(dim, True)\n            total_zeros = slice_zero_count.data.sum()\n            if total_zeros == 0:\n                grad_input = grad_output.mul(output).expand_as(input).div(input)\n            else:\n                grad_input = safe_zeros_backward(input, dim)\n\n            return grad_input, None, None\n\n\nclass Mean(Function):\n\n    @staticmethod\n    def forward(ctx, input, dim=None, keepdim=None):\n        ctx.dim = dim\n        ctx.keepdim = False if keepdim is None else keepdim\n        ctx.input_size = input.size()\n        if dim is None:\n            return input.new((input.mean(),))\n        else:\n            if keepdim is not None:\n                return input.mean(dim, keepdim=keepdim)\n            else:\n                return input.mean(dim)\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        if ctx.dim is None:\n            grad_input_val = grad_output / reduce(lambda x, y: x * y, ctx.input_size, 1)\n            return grad_input_val.expand(ctx.input_size), None, None\n        else:\n            if ctx.keepdim is False and len(ctx.input_size) != 1:\n                grad_output = grad_output.unsqueeze(ctx.dim)\n\n            repeats = [1 for _ in ctx.input_size]\n            dim_size = ctx.input_size[ctx.dim]\n            repeats[ctx.dim] = dim_size\n            return grad_output.repeat(*repeats).div_(dim_size), None, None\n\n\nclass _SelectionFunction(Function):\n    has_all_reduce = True\n    # additional_args is prepended before dim when calling the tensor\n    # function. It's a no-op for subclasses other than kthvalue.\n    # kthvalue not only requires us to pass a dim, but also precede it with k.\n\n    @classmethod\n    def forward(cls, ctx, input, dim=None, keepdim=None, additional_args=tuple()):\n        fn = getattr(input, cls.__name__.lower())\n        ctx.dim = dim\n        ctx.keepdim = False if keepdim is None else keepdim\n        ctx.additional_args = additional_args\n        ctx.input_size = input.size()\n        if ctx.dim is None and cls.has_all_reduce:\n            value = fn(*additional_args)\n            ctx.indices_tuple = tuple(input.eq(value).nonzero()[0])\n            return input.new((value,))\n        else:\n            if ctx.dim is None:\n                dim = input.dim() - 1\n            else:\n                dim = ctx.dim\n            args = (dim,)\n            if additional_args:\n                args = additional_args + args\n            if keepdim is not None:\n                output, indices = fn(*args, keepdim=keepdim)\n            else:\n                output, indices = fn(*args)\n            ctx.save_for_backward(indices)\n            ctx.mark_non_differentiable(indices)\n            return output, indices\n\n    @classmethod\n    def backward(cls, ctx, grad_output, grad_indices=None):\n        grad_input = Variable(grad_output.data.new(*ctx.input_size).zero_())\n        if ctx.dim is None and cls.has_all_reduce:\n            grad_input[ctx.indices_tuple] = grad_output\n        else:\n            if ctx.dim is None:\n                dim = len(ctx.input_size) - 1\n            else:\n                dim = ctx.dim\n\n            indices, = ctx.saved_variables\n            if ctx.keepdim is False and len(ctx.input_size) != 1:\n                grad_output = grad_output.unsqueeze(dim)\n                grad_indices = grad_indices.unsqueeze(dim)\n                indices = indices.unsqueeze(dim)\n\n            grad_input.scatter_(dim, indices, grad_output)\n        return grad_input, None, None, None\n\n\nclass Max(_SelectionFunction):\n    pass\n\n\nclass Min(_SelectionFunction):\n    pass\n\n\nclass Mode(_SelectionFunction):\n    has_all_reduce = False\n\n\nclass Median(_SelectionFunction):\n    pass\n\n\nclass Kthvalue(_SelectionFunction):\n    has_all_reduce = False\n\n    @classmethod\n    def forward(cls, ctx, input, k, dim=None, keepdim=None):\n        return super(Kthvalue, cls).forward(ctx, input, dim, keepdim, (k,))\n\n\nclass Norm(Function):\n\n    @staticmethod\n    def forward(ctx, input, p=2, dim=None, keepdim=None):\n        ctx.p = p\n        ctx.dim = dim\n        ctx.keepdim = False if keepdim is None else keepdim\n\n        if dim is None:\n            norm = input.norm(p)\n            output = input.new((norm,))\n        else:\n            if keepdim is not None:\n                output = input.norm(p, dim, keepdim=keepdim)\n            else:\n                output = input.norm(p, dim)\n        ctx.save_for_backward(input, output)\n        return output\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        input, output = ctx.saved_variables\n        if ctx.dim is not None and ctx.keepdim is False and input.dim() != 1:\n            grad_output = grad_output.unsqueeze(ctx.dim)\n            output = output.unsqueeze(ctx.dim)\n\n        if ctx.p == 2:\n            grad_input = input.mul(grad_output).div(output)\n        else:\n            input_pow = input.abs().pow(ctx.p - 2)\n            output_pow = output.pow(ctx.p - 1)\n            grad_input = input.mul(input_pow).mul(grad_output).div(output_pow)\n\n        # Special case at 0 where we return a subgradient containing 0\n        grad_input.masked_fill_(output == 0, 0)\n\n        return grad_input, None, None, None\n\n\n# TODO: renorm\n# TODO: std\n# TODO: var\n", "path": "torch/autograd/_functions/reduce.py"}]}