Repository: pytorch/pytorch

Issue:
the derivative for 'cudnn_batch_norm_backward save_mean' is not implemented
```
loss_gp = ..
loss_gp.backward()  # error
```

Having this error when BatchNorm + wgan_gp (second order) is used.
It's installed from master fec3d4.



Discussion:
This is almost definitely my fault. I'll look into it on Tuesday.
@elbaro Would it be possible to get more detailed reproduction instructions?
Never mind, it's pretty easy to repro:

```
    def test_batchnorm_double_backward(self):
        bn = nn.BatchNorm1d(5).cuda()
        v = Variable(torch.randn(5, 5), requires_grad=True).cuda()
        output = bn(v)
        grad, = torch.autograd.grad(outputs=output.sum(), inputs=v, create_graph=True, retain_graph=True, only_inputs=True)
        grad.sum().backward()
```
Root causes:

- I did a hack to work around lack of ability to return non-requires grad outputs from ATen functions, but apparently this hack doesn't work. We need a way of returning variables from an ATen function such that they unconditionally don't require grad. @colesbury suggested that we add a new special form to `derivatives.yaml` to trigger this but I am actually not too sure we can actually make this work. In any case, @colesbury has some PRs in progress on this code so I won't be able to fix this today.
- Why didn't I catch the hack failure? We are apparently not running any NewModuleTests double-backwards with CUDA. Oops. I'm going to try to fix this today.

To work around the problem:
- Turn off cuDNN with `torch.backends.cudnn.enabled = False`, OR
- Revert to PyTorch d7e6ede78465ceddf574840eb844f08a9e870557 or earlier

Resolving patch:
diff --git a/tools/autograd/load_derivatives.py b/tools/autograd/load_derivatives.py
--- a/tools/autograd/load_derivatives.py
+++ b/tools/autograd/load_derivatives.py
@@ -92,7 +92,61 @@ def lookup_pred(pred, xs):
         """Return the index of the first element of xs matching pred."""
         return next((i, x) for i, x in enumerate(xs) if pred(x))
 
-    def set_up_derivatives(defn, declaration):
+    def check_grad_usage(defn_name, declaration, derivatives):
+        """
+        Check for some subtle mistakes one might make when writing gradients.
+        These mistakes will compile, but will be latent until a function is
+        used with double backwards.
+        """
+
+        used_grad = 0
+        used_grads = 0
+        fully_implemented = True
+        used_grads_indices = []
+        for d in derivatives:
+            formula = d['formula']
+            used_grad += len(re.findall(IDENT_REGEX.format('grad'), formula))
+            used_grads += len(re.findall(IDENT_REGEX.format('grads'), formula))
+            fully_implemented = \
+                fully_implemented and \
+                not re.search(IDENT_REGEX.format('not_implemented'), formula)
+            used_grads_indices.extend(used_gradient_indices(formula))
+        assert used_grads >= len(used_grads_indices)
+        only_used_grads_indices = used_grads == len(used_grads_indices)
+
+        if used_grad and used_grads:
+            raise RuntimeError("Derivative definition of {} in derivatives.yaml illegally "
+                               "mixes use of 'grad' and 'grads'. Consider replacing "
+                               "occurrences of 'grad' with 'grads[0]'".format(defn_name))
+
+        if only_used_grads_indices and set(used_grads_indices) == {0}:
+            raise RuntimeError("Derivative definition of {} in derivatives.yaml solely "
+                               "refers to 'grads[0]'.  If the first output is indeed the "
+                               "only differentiable output, replace 'grads[0]' with 'grad'; "
+                               "otherwise, there is a likely error in your derivatives "
+                               "declaration.".format(defn_name))
+
+        # DO NOT comment out this test!  Code generation will probably work with
+        # this test commented out, but if you ever pass a non-differentiable
+        # argument to an autograd function (e.g., a backwards function which
+        # has double backwards implemented, as was the case in #4422) your code
+        # will fail when you ever actually try to differentiate with it.
+        #
+        # NB: I had to make it not complain if both 'grads' and 'grad' are never
+        # used, because we have some silly zeros_like() gradients for inplace
+        # comparison tests.
+        if fully_implemented and not used_grad and used_grads and only_used_grads_indices and \
+           set(used_grads_indices) != set(range(len(declaration['returns']))):
+            raise RuntimeError("Derivative definition of {} in derivatives.yaml does "
+                               "not refer to the gradients of all of its outputs.  Either "
+                               "the derivatives declaration is wrong, OR you have some "
+                               "non-differentiable outputs.  If you have a single "
+                               "differentiable output, make it the first output in ATen "
+                               "and reference its gradient with 'grad'; otherwise, you "
+                               "have hit a case which is unsupported by the codegen, "
+                               "see #4567.".format(defn_name))
+
+    def set_up_derivatives(defn_name, defn, declaration):
         # First, let us determine the set of inputs for which gradients
         # were specified in declarations.  We'll use this in layout
         # computation.
@@ -129,6 +183,9 @@ def set_up_derivatives(defn, declaration):
                 args.append(name)
             derivatives.append(create_derivative(declaration, formula, output_indices, args))
 
+        # Test to see if the use of 'grads' makes sense.
+        check_grad_usage(defn_name, declaration, derivatives)
+
         return derivatives, num_inputs
 
     def unzip(xs):
@@ -168,7 +225,7 @@ def unzip(xs):
                                'Declarations.yaml ({})'
                                .format(i, defn_name, x, y))
 
-    derivatives, num_inputs = set_up_derivatives(defn, canonical)
+    derivatives, num_inputs = set_up_derivatives(defn_name, defn, canonical)
     return create_autograd_function(defn_name, derivatives, num_inputs, signature)
 
 
@@ -204,10 +261,6 @@ def used_gradient_indices(formula):
     """Determine a list of gradient indices (the i in grads[i]) that
     are used by the formula.
 
-    NB: references to 'grad' don't count as a gradient index (technically,
-    we should return 0 in this case, but it doesn't matter this can never
-    be out of bounds
-
     >>> used_gradient_indices("foo(grads[0], grads[1])")
     [0, 1]
     """


Test patch:
diff --git a/test/common_nn.py b/test/common_nn.py
--- a/test/common_nn.py
+++ b/test/common_nn.py
@@ -834,17 +834,19 @@ def test_cuda(self, test_case):
 
             # Run double-backwards on CPU and GPU and compare results
             if self.check_gradgrad and not self.FIXME_no_cuda_gradgrad_comparison:
-                cpu_output_t = cpu_output.data if isinstance(cpu_output, Variable) else cpu_output
-                cpu_gradOutput = Variable(cpu_output_t.clone().normal_(), requires_grad=True)
-                gpu_gradOutput = Variable(cpu_gradOutput.type('torch.cuda.FloatTensor').data, requires_grad=True)
+                cpu_output = cpu_module(cpu_input)
+                gpu_output = gpu_module(gpu_input)
+
+                cpu_gradOutput = Variable(cpu_output.data.clone().normal_(), requires_grad=True)
+                gpu_gradOutput = Variable(cpu_gradOutput.data.type_as(gpu_output), requires_grad=True)
 
                 cpu_gradInputs = torch.autograd.grad(
-                    cpu_module(cpu_input),
+                    cpu_output,
                     (cpu_input,) + tuple(cpu_module.parameters()),
                     cpu_gradOutput,
                     create_graph=True)
                 gpu_gradInputs = torch.autograd.grad(
-                    gpu_module(gpu_input),
+                    gpu_output,
                     (gpu_input,) + tuple(gpu_module.parameters()),
                     gpu_gradOutput,
                     create_graph=True)
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -4530,7 +4530,6 @@ def smoothl1loss_no_reduce_test():
         cudnn=True,
         check_eval=True,
         desc='affine',
-        FIXME_no_cuda_gradgrad_comparison=True,  # See #4422
     ),
     dict(
         module_name='BatchNorm1d',
@@ -4539,7 +4538,6 @@ def smoothl1loss_no_reduce_test():
         cudnn=True,
         check_eval=True,
         desc='3d_input',
-        FIXME_no_cuda_gradgrad_comparison=True,  # See #4422
     ),
     dict(
         module_name='BatchNorm1d',
@@ -4548,7 +4546,6 @@ def smoothl1loss_no_reduce_test():
         cudnn=True,
         check_eval=True,
         desc='not_affine',
-        FIXME_no_cuda_gradgrad_comparison=True,  # See #4422
     ),
     dict(
         module_name='BatchNorm1d',
@@ -4557,7 +4554,6 @@ def smoothl1loss_no_reduce_test():
         cudnn=True,
         check_eval=True,
         desc='3d_input_not_affine',
-        FIXME_no_cuda_gradgrad_comparison=True,  # See #4422
     ),
     dict(
         module_name='BatchNorm2d',
@@ -4565,7 +4561,6 @@ def smoothl1loss_no_reduce_test():
         input_size=(2, 3, 6, 6),
         cudnn=True,
         check_eval=True,
-        FIXME_no_cuda_gradgrad_comparison=True,  # See #4422
     ),
     dict(
         module_name='BatchNorm2d',
@@ -4574,7 +4569,6 @@ def smoothl1loss_no_reduce_test():
         cudnn=True,
         check_eval=True,
         desc='momentum',
-        FIXME_no_cuda_gradgrad_comparison=True,  # See #4422
     ),
     dict(
         module_name='BatchNorm2d',
@@ -4583,7 +4577,6 @@ def smoothl1loss_no_reduce_test():
         cudnn=True,
         check_eval=True,
         desc='not_affine',
-        FIXME_no_cuda_gradgrad_comparison=True,  # See #4422
     ),
     dict(
         module_name='BatchNorm3d',
@@ -4591,7 +4584,6 @@ def smoothl1loss_no_reduce_test():
         input_size=(2, 3, 4, 4, 4),
         cudnn=True,
         check_eval=True,
-        FIXME_no_cuda_gradgrad_comparison=True,  # See #4422
     ),
     dict(
         module_name='BatchNorm3d',
@@ -4600,7 +4592,6 @@ def smoothl1loss_no_reduce_test():
         cudnn=True,
         check_eval=True,
         desc='momentum',
-        FIXME_no_cuda_gradgrad_comparison=True,  # See #4422
     ),
     dict(
         module_name='BatchNorm3d',
@@ -4609,7 +4600,6 @@ def smoothl1loss_no_reduce_test():
         cudnn=True,
         check_eval=True,
         desc='not_affine',
-        FIXME_no_cuda_gradgrad_comparison=True,  # See #4422
     ),
     dict(
         module_name='Conv1d',