Cleanup jvp testing (#890)

- deleted test_vmapjvp (test_vmapjvpall is a superset of test_vmapjvp)
- added additional_op_db to things tested by test_vmapjvpall
- did some accounting in discover_coverage

Author: zou3519

test/discover_coverage.py

@@ -628,8 +628,6 @@ class Support(enum.Enum):
     'nn.functional.rrelu',  # not actually problem, randomness testing artifact
     'normal',  # not actually problem, randomness testing artifact
     'bernoulli',  # not actually problem, randomness testing artifact
-    'nn.functional.embedding',  # max_norm causes testing to be weird
-    # 'multinomial',
 }
 
 
@@ -746,8 +744,12 @@ def _supports_vmapjvp_base(self, test):
             'nn.functional.batch_norm',  # testing problem
             'normal',  # not actually problem, randomness testing artifact
             'bernoulli',  # not actually problem, randomness testing artifact
-            'dropout2d',  # not actually problem, randomness testing artifact
-            'dropout',  # not actually problem, randomness testing artifact
+            'nn.functional.dropout2d',  # not actually problem, randomness testing artifact
+            'nn.functional.dropout',  # not actually problem, randomness testing artifact
+            # Not a problem.
+            # It's just that the max_norm testing mutates inputs...
+            # (we have our own functorch variant of the OpInfo without max_norm)
+            'nn.functional.embedding',
         }
         if self.name in VMAPJVP_EXEMPTIONS:
             return Support.YES

test/functorch_additional_op_db.py

@@ -181,6 +181,8 @@ def generator():
         op=lambda weight, idx, **kwargs: torch.nn.functional.embedding(idx, weight, **kwargs),
         dtypes=floating_types_and(torch.bfloat16, torch.float16),
         sample_inputs_func=sample_inputs_embedding,
+        supports_forward_ad=True,
+        supports_fwgrad_bwgrad=True,
         supports_out=False,
     ))
 

test/test_ops.py

@@ -648,92 +648,6 @@ def test_vmapvjp(self, device, dtype, op):
             for loop_out, batched_out in get_fallback_and_vmap_exhaustive(fn, args, {}, opinfo=op):
                 self.assertEqual(loop_out, batched_out)
 
-    # There are several variations we care about
-    # 1) primal batched (TODO)
-    # 2) tangent batched (batched grads) <--
-    # 3) both batched (TODO)
-    # The below tests (2) only.
-    @ops(functorch_lagging_op_db, allowed_dtypes=(torch.float,))
-    @toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1e-04)})
-    @skipOps('TestOperators', 'test_vmapjvp', {
-        skip('nn.functional.dropout'),  # randomness
-        skip('nn.functional.rrelu'),  # randomness
-        skip('nn.functional.fractional_max_pool2d'),  # randomness
-        skip('nn.functional.fractional_max_pool3d'),  # randomness
-        skip('bernoulli', ''),  # randomness
-        skip('nn.functional.max_pool1d'),  # fails on cpu, runs on cuda
-
-        # TODO: fails in core due to in-place batched nto non-batched
-        # but fails here for a different reason
-        xfail('linalg.householder_product'),
-
-        # Try to in-place batched tensor into non-batched tensor
-        xfail('matrix_exp'),
-
-        # Apprently these support forward AD, but we get "Trying to use forward AD..."
-        # These are cases where OpInfo has supports_forward_ad=True, but disables
-        # the test
-        xfail('var_mean'),
-        xfail('std_mean'),
-
-        # RuntimeError: expand: the number of sizes provided (1) must be greater or
-        # equal to the number of dimensions in the tensor (2)
-        xfail('nanquantile'),
-        xfail('quantile'),
-
-        # Not implemented
-        xfail('scatter'),
-
-        # =============================================
-        # NB: The above failures also fail in PyTorch core.
-        #     The failures below only fail in functorch
-        # =============================================
-
-        # Composite ops that do bad things. Need to be fixed in PyTorch core.
-        # RuntimeError: Cannot access data pointer of Tensor that doesn't have storage
-        xfail('tensor_split'),
-
-        # Causing multiple forward mode AD issues, needs investigation
-        xfail('nn.functional.batch_norm'),
-        xfail('nn.functional.batch_norm', 'without_cudnn', device_type='cuda'),
-
-        skip('nn.functional.feature_alpha_dropout', 'with_train'),
-        skip('pca_lowrank', ''),
-        skip('nn.functional.dropout2d', ''),
-        skip('nn.functional.feature_alpha_dropout', 'without_train'),
-        skip('svd_lowrank', ''),
-        xfail('nn.functional.soft_margin_loss', ''),
-        xfail('stft'),  # something weird is happening with shapes
-
-        xfail('double'),  # required rank 4 tensor to use channels_last format
-
-        # BUG: runs and produces numerical differences
-        skip('nn.functional.max_unpool1d', device_type='cpu'),  # fails everywhere except on mac
-        skip('nn.functional.max_unpool2d'),  # fails everywhere except on mac
-        skip('nn.functional.max_unpool3d'),  # fails everywhere except on mac
-
-        xfail('put'),  # calls put_ during vmap with only vmaps over other, not self
-    })
-    def test_vmapjvp(self, device, dtype, op):
-        if is_inplace(op, op.get_op()):
-            # TODO: test in-place
-            self.skipTest("Skipped! NYI: inplace-testing not supported.")
-            return
-
-        samples = op.sample_inputs(device, dtype, requires_grad=False)
-
-        if not op.supports_forward_ad:
-            self.skipTest("Skipped! Forward AD not supported.")
-            return
-
-        for sample in samples:
-            arg_values = [sample.input] + list(sample.args)
-            kwarg_values = sample.kwargs
-            args = tuple([*arg_values, *kwarg_values])
-            fn, args = get_jvp_variant(op, sample)
-            for loop_out, batched_out in get_fallback_and_vmap_exhaustive(fn, args, {}, opinfo=op, bdims=(0,)):
-                self.assertEqual(loop_out, batched_out)
-
     vmapjvpall_fail = {
         # The following are expected (not a bug)
         skip('bernoulli', ''),  # randomness
@@ -757,7 +671,8 @@ def test_vmapjvp(self, device, dtype, op):
 
         # Not actually a problem: embedding with max_norm mutates the weight
         # and causes different runs to produce different results.
-        xfail('nn.functional.embedding', ''),
+        # skip because this is flaky depending on what the max_norm is!
+        skip('nn.functional.embedding', ''),
         xfail('nn.functional.soft_margin_loss', ''),
         xfail('nn.functional.binary_cross_entropy_with_logits', ''),
         xfail('linalg.householder_product'),
@@ -788,7 +703,7 @@ def test_vmapjvp(self, device, dtype, op):
         xfail('nn.functional.prelu'),  # Call Tensor.as_strided
     }
 
-    @ops(functorch_lagging_op_db, allowed_dtypes=(torch.float,))
+    @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
     @opsToleranceOverride('TestOperators', 'test_vmapjvpall', (
         tol1('nn.functional.conv_transpose3d',
              {torch.float32: tol(atol=2e-04, rtol=9e-3)}, device_type='cuda'),
@@ -818,7 +733,7 @@ def test_vmapjvpall(self, device, dtype, op):
             for loop_out, batched_out in get_fallback_and_vmap_exhaustive(fn, args, {}, opinfo=op):
                 self.assertEqual(loop_out, batched_out)
 
-    @ops(functorch_lagging_op_db, allowed_dtypes=(torch.float,))
+    @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
     @skipOps('TestOperators', 'test_vmapjvpall_has_batch_rule', vmapjvpall_fail.union({
         xfail('linalg.solve_triangular'),
         xfail('nn.functional.huber_loss'),