Allow people to arbitrarily add dispatch keys between DynamicLayer{Front,Back} (#843)

Fixes #842

The Diagnosis
=============
As Brian pointed out:

For jvp(sub, ...), the chain of dispatch should be:

```
DynamicLayerFrontMode -> at::sub autograd kernel -> DynamicLayerBackMode
```

Instead, what we're doing today is
```
JVP dynamic layer -> at::sub autograd kernel -> at::sub zero_kernel
```
(the zero_tensor kernel errors out, because the inputs are
BatchedTensorImpl objects)

The Problem
=============

functorch's behavior on dispatch keys between DynamicLayerFrontMode and
DynamicLayerBack mode should be:
- upon entering a dynamic layer (aka Interpreter), we zero out all
dispatch keys* between FrontMode and BackMode
- then, the dynamic layer (aka Interpreter) decides to re-enable some
dispatch keys. For example, JVPInterpreter decides to re-enable the
autograd keys
- next, we do a dispatcher call, which will end up hitting one of the
Autograd keys (in the JVPInterpreter case).

The bug is that functorch has a hardcoded list of dispatch keys that it
zeros out. This list does not include ZeroTensor, because before
pytorch/pytorch#77132, the ZeroTensor key was
not between DynamicLayer{Front,Back}Mode.

*There is an exception for autocast and vmapmode, described in the next section.

The Solution
============

Change functorch to programmatically zero out keys between
DynamicLayerBackMode and DynamicLayerFrontMode, with the exception of
Autocast and VmapMode.

This means that in the future, if someone adds a dispatch key between
DynamicLayerBackMode and DynamicLayerFrontMode, we will (probably) be
handling it "correctly": the model for dispatch is:
- [functorch] -> [regular pytorch dispatcher]
- a key like ZeroTensor gets handled in the [regular pytorch dispatcher]
section.
- functorch transforms get handled in the [functorch] section.

We do not change the autocast <-> functorch interaction in this PR
(i.e. functorch does not zero it out) because I'm not sure what the
correct thing to do here is.

We do not change how kVmapMode works because... it needs to be active
to ban random operations in transforms later down the line :/

Test Plan
============
Wait for tests

Author: zou3519

functorch/csrc/Interpreter.cpp

@@ -7,16 +7,30 @@
 
 namespace at { namespace functorch {
 
-constexpr DispatchKeySet all_dynlayer_keyset = DispatchKeySet({
-  kDynamicLayerFrontModeKey,
-  kDynamicLayerBackModeKey,
-  kGradWrapperKey,
-  DispatchKey::Functionalize,
-  // DispatchKey::Batched,
-  kBatchedKey,
-  DispatchKey::PythonTLSSnapshot,
-  DispatchKey::ADInplaceOrView
-}) | autograd_dispatch_keyset;
+static DispatchKeySet get_all_dynlayer_keyset() {
+  // NB: FULL_AFTER does not include the dispatch key
+
+  // "all dispatch keys between DynamicLayer{Front, Back}Mode, inclusive"
+  auto result =
+    DispatchKeySet(DispatchKeySet::FULL_AFTER, kDynamicLayerFrontModeKey) -
+    DispatchKeySet(DispatchKeySet::FULL_AFTER, kDynamicLayerBackModeKey);
+  result = result | DispatchKeySet({kDynamicLayerFrontModeKey});
+
+  // Hack: don't handle the autocast dispatch keys. Their interaction with functorch
+  // is weird.
+  result = result - autocast_dispatch_keyset;
+
+  // Hack: don't handle kVmapModeKey. We need a better way of modeling this.
+  // In e.g. grad(vmap(f)), kVmapModeKey makes it so that all random operations,
+  // even after we are done handling the vmap layer, error out.
+  result = result.remove(kVmapModeKey);
+
+  return result;
+}
+
+// TODO: This should be constexpr, but there are some methods
+// of DispatchKeySet that haven't been marked constexpr yet.
+static DispatchKeySet all_dynlayer_keyset = get_all_dynlayer_keyset();
 
 static DispatchKeySet keysForEnteringDynamicLayer(TransformType key) {
   if (key == TransformType::Vmap) {

test/test_eager_transforms.py

@@ -1999,6 +1999,18 @@ def backward(ctx, gx):
         gx, = torch.autograd.grad(y, x)
         self.assertEqual(gx, x.cos())
 
+    def test_zerotensor_vmapjvp_interaction(self, device):
+        dummy = torch.ones(4, 1)
+        x = torch.randn(4, 2)
+        x_tangent = torch.randn(2)
+
+        def push_jvp(dummy, x):
+            result = jvp(torch.cov, (x,), (x_tangent,))
+            return result
+
+        # Should not error
+        vmap(vmap(push_jvp, (0, None)))(dummy, x)
+
 
 class TestCustomFunction(TestCase):
     @unittest.skipIf(IS_WINDOWS, "Prototype of custom_vjp doesn't link on windows")