"set_inplace_requires_grad_allowed" should be a context manager

functorch/_src/eager_transforms.py

@@ -33,20 +33,27 @@
     _assert_wrapped_functional,
     _propagate_functional_input_mutation,
     set_inplace_requires_grad_allowed,
+    get_inplace_requires_grad_allowed,
 )
 
 argnums_t = Union[int, Tuple[int, ...]]
 
 
+@contextlib.contextmanager
+def enable_inplace_requires_grad(enabled=True):
+    prev_state = get_inplace_requires_grad_allowed()
+    set_inplace_requires_grad_allowed(enabled)
+    try:
+        yield
+    finally:
+        set_inplace_requires_grad_allowed(prev_state)
+
+
 def _create_differentiable(inps, level=None):
     def create_differentiable(x):
         if isinstance(x, torch.Tensor):
-            try:
-                set_inplace_requires_grad_allowed(True)
+            with enable_inplace_requires_grad():
                 return x.requires_grad_()
-            finally:
-                set_inplace_requires_grad_allowed(False)
-
         raise ValueError(f'Thing passed to transform API must be Tensor, '
                          f'got {type(x)}')
     return tree_map(create_differentiable, inps)

functorch/csrc/DynamicLayer.cpp

@@ -140,6 +140,11 @@ void setInplaceRequiresGradAllowed(bool allowed) {
   functorch_tls->allow_inplace_requires_grad_ = allowed;
 }
 
+bool getInplaceRequiresGradAllowed() {
+  auto* functorch_tls = getRawFunctorchTLS();
+  return functorch_tls->allow_inplace_requires_grad_;
+}
+
 
 static std::vector<DynamicLayer>& dynamicLayerStackAccessor() {
   return getRawFunctorchTLS()->dynamicLayerStack;

functorch/csrc/DynamicLayer.h

@@ -86,6 +86,7 @@ std::ostream& operator<<(std::ostream& os, const DynamicLayer& layer);
 std::ostream& operator<<(std::ostream& os, const std::vector<DynamicLayer>& dynamicLayerStack);
 
 void setInplaceRequiresGradAllowed(bool allowed);
+bool getInplaceRequiresGradAllowed();
 
 
 }

functorch/csrc/init.cpp

@@ -380,6 +380,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("_set_vmap_fallback_enabled", &at::functorch::setVmapFallbackEnabled);
   m.def("_is_vmap_fallback_enabled", &at::functorch::isVmapFallbackEnabled);
   m.def("set_inplace_requires_grad_allowed", &at::functorch::setInplaceRequiresGradAllowed);
+  m.def("get_inplace_requires_grad_allowed", &at::functorch::getInplaceRequiresGradAllowed);
   m.def("dlevel", &at::functorch::dlevel, "dlevel");
   m.def("dump_tensor", &at::functorch::dump_tensor, "dump_tensor");
   m.def("reshape_dim_into", &at::functorch::reshape_dim_into);