[functorch] make jvpvjp actually run all the tests for decompositions (pytorch/functorch#792)

Author: Samantha Andow

functorch/functorch/_src/eager_transforms.py

@@ -1302,6 +1302,7 @@ def get_function_def(sig):
 
 _register_jit_decomposition(torch.ops.aten.trace.default)
 _register_jit_decomposition(torch.ops.aten.nll_loss_backward.default)
+_register_jit_decomposition(torch.ops.aten.nll_loss2d_backward.default)
 _register_jit_decomposition(torch.ops.aten.mse_loss_backward.default)
 _register_jit_decomposition(torch.ops.aten.l1_loss_backward.default)
 _register_jit_decomposition(torch.ops.aten._log_softmax_backward_data.default)

functorch/test/test_ops.py

@@ -1359,10 +1359,9 @@ def get_vjp(primals, cotangents):
                 if op.name == 'nn.functional.binary_cross_entropy':  # reverse second derivative wrt target not defined
                     in_dims = 1
                 compare_jacobians(primals, cotangents, in_dims)
-                return
-
-            expected = reference(primals, cotangents, primals_tangents, cotangents_tangents)
-            self.assertEqual(result, expected)
+            else:
+                expected = reference(primals, cotangents, primals_tangents, cotangents_tangents)
+                self.assertEqual(result, expected)
 
     @ops(filter(lambda op: op.name == "nn.functional.group_norm", functorch_lagging_op_db + additional_op_db),
          allowed_dtypes=(torch.float32, torch.double))  # TODO: generalize