Apply functorch-side fix for PythonTLSSnapshot change (#633)

We add PythonTLSSnapshot to the local exclude while functorch is active
and remove it when we leave functorch. This is a quick fix; there's
something nicer we can do (which is refactor DynamicLayer to not have so
many RAII guards that modify the thread-local dispatcher state)

Author: zou3519

functorch/csrc/DynamicLayer.cpp

@@ -26,6 +26,7 @@ constexpr DispatchKeySet all_dynlayer_keyset = DispatchKeySet({
   DispatchKey::Functionalize,
   // DispatchKey::Batched,
   kBatchedKey,
+  DispatchKey::PythonTLSSnapshot,
   DispatchKey::ADInplaceOrView
 }) | autograd_dispatch_keyset;
 
@@ -495,6 +496,9 @@ void dynamicLayerFrontFallback(const c10::OperatorHandle& op, torch::jit::Stack*
   if (dynamicLayerStack.size() == 0) {
     sanityCheckStack(op, stack);
     c10::impl::ExcludeDispatchKeyGuard guard(all_dynlayer_keyset);
+    auto local_keyset = c10::impl::tls_local_dispatch_key_set();
+    local_keyset.excluded_ = local_keyset.excluded_.remove(DispatchKey::PythonTLSSnapshot);
+    c10::impl::ForceDispatchKeyGuard guard2(local_keyset);
     op.callBoxed(stack);
     return;
   }

test/test_eager_transforms.py

@@ -23,7 +23,7 @@
 from functorch import (
     grad, vjp, vmap, jacrev, jacfwd, grad_and_value, hessian,
     jvp, make_functional, make_functional_with_buffers,
-    combine_state_for_ensemble,
+    combine_state_for_ensemble, make_fx
 )
 from functorch._src.make_functional import (
     functional_init, functional_init_with_buffers,
@@ -2112,6 +2112,35 @@ def test_vjp_vjp(self, device):
         y = vjp_fn(x)[0]
         # Honestly IDK what the result here is... but at least it runs
 
+    def test_make_fx_vmap(self, device):
+        def f(x):
+            return torch.sin(x)
+        inp = torch.randn(5, 3)
+        f = vmap(f)
+        fx_f = make_fx(f)(inp)
+        new_inp = torch.randn(5, 3)
+        self.assertEqual(fx_f(new_inp), f(new_inp))
+
+    def test_make_fx_jacrev(self, device):
+        def f(x):
+            return x.sin().sum()
+        inp = torch.randn(3)
+        f = jacrev(jacrev(f))
+        fx_f = make_fx(f)(inp)
+        new_inp = torch.randn(3)
+        self.assertEqual(fx_f(new_inp), f(new_inp))
+
+    def test_make_fx_vjp(self, device):
+        def f(x):
+            return torch.sin(x).sum()
+
+        primals = torch.randn(3)
+        _, vjp_fn = vjp(f, primals)
+        cotangent = torch.randn(())
+        fx_f = make_fx(vjp_fn)(cotangent, True, True)
+        new_cotangent = torch.randn(())
+        self.assertEqual(fx_f(new_cotangent, True, True), vjp_fn(new_cotangent))
+
 
 class TestMakeFunctional(TestCase):
     def test_parameter_tying(self):