Move torch/autograd/grad_mode.pyi stubs inline (pytorch#43415)

Summary:
- Add `torch._C` bindings from `torch/csrc/autograd/init.cpp`
- Renamed `torch._C.set_grad_enabled` to `torch._C._set_grad_enabled`
  so it doesn't conflict with torch.set_grad_enabled anymore

This is a continuation of pytorchgh-38201. All I did was resolve merge conflicts and finish the annotation of `_DecoratorContextManager.__call__` that ezyang started in the first commit.

~Reverts commit b5cd3a8, which was only motivated by not having `typing_extensions` available.~ (JIT can't be made to understand `Literal[False]`, so keep as is).

Pull Request resolved: pytorch#43415

Reviewed By: ngimel

Differential Revision: D23301168

Pulled By: malfet

fbshipit-source-id: cb5290f2e556b4036592655b9fe54564cbb036f6

Author: rgommers

torch/_C/__init__.pyi.in

@@ -217,6 +217,17 @@ has_cudnn: _bool
 _GLIBCXX_USE_CXX11_ABI: _bool
 default_generator: Generator
 
+# Defined in torch/csrc/autograd/init.cpp
+def _set_grad_enabled(enabled: _bool) -> None: ...
+def is_grad_enabled() -> _bool: ...
+def set_autocast_enabled(enabled: _bool) -> None: ...
+def is_autocast_enabled() -> _bool: ...
+def clear_autocast_cache() -> None: ...
+def autocast_increment_nesting() -> None: ...
+def autocast_decrement_nesting() -> None: ...
+def set_anomaly_enabled(enabled: _bool) -> None: ...
+def is_anomaly_enabled() -> _bool: ...
+
 # Defined in torch/csrc/jit/python/script_init.cpp
 class FileCheck(object):
     # TODO

torch/autograd/grad_mode.py

@@ -1,20 +1,30 @@
 import torch
 import functools
 import inspect
-from typing import Any
+from typing import Any, Callable, TypeVar, cast
+
+
+__all__ = ['no_grad', 'enable_grad', 'set_grad_enabled']
+
+
+# Used for annotating the decorator usage of 'no_grad' and 'enable_grad'.
+# See https://mypy.readthedocs.io/en/latest/generics.html#declaring-decorators
+FuncType = Callable[..., Any]
+F = TypeVar('F', bound=FuncType)
+
 
 class _DecoratorContextManager:
     """Allow a context manager to be used as a decorator"""
 
-    def __call__(self, func):
+    def __call__(self, func: F) -> F:
         if inspect.isgeneratorfunction(func):
             return self._wrap_generator(func)
 
         @functools.wraps(func)
         def decorate_context(*args, **kwargs):
             with self:
                 return func(*args, **kwargs)
-        return decorate_context
+        return cast(F, decorate_context)
 
     def _wrap_generator(self, func):
         """Wrap each generator invocation with the context manager"""
@@ -30,6 +40,12 @@ def generator_context(*args, **kwargs):
                     break
         return generator_context
 
+    def __enter__(self) -> None:
+        raise NotImplementedError
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        raise NotImplementedError
+
 
 class no_grad(_DecoratorContextManager):
     r"""Context-manager that disabled gradient calculation.
@@ -70,7 +86,7 @@ def __enter__(self):
         self.prev = torch.is_grad_enabled()
         torch.set_grad_enabled(False)
 
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         torch.set_grad_enabled(self.prev)
 
 
@@ -105,12 +121,12 @@ class enable_grad(_DecoratorContextManager):
         True
 
     """
-    def __enter__(self):
+    def __enter__(self) -> None:
         self.prev = torch.is_grad_enabled()
-        torch._C.set_grad_enabled(True)
+        torch._C._set_grad_enabled(True)
 
-    def __exit__(self, *args):
-        torch.set_grad_enabled(self.prev)
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        torch._C._set_grad_enabled(self.prev)
 
 
 class set_grad_enabled(object):
@@ -147,12 +163,12 @@ class set_grad_enabled(object):
 
     """
 
-    def __init__(self, mode):
+    def __init__(self, mode: bool) -> None:
         self.prev = torch.is_grad_enabled()
-        torch._C.set_grad_enabled(mode)
+        torch._C._set_grad_enabled(mode)
 
-    def __enter__(self):
+    def __enter__(self) -> None:
         pass
 
-    def __exit__(self, *args):
-        torch.set_grad_enabled(self.prev)
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        torch._C._set_grad_enabled(self.prev)

torch/autograd/grad_mode.pyi

@@ -163,7 +163,7 @@ static PyObject * is_anomaly_mode_enabled(PyObject* _unused, PyObject *arg) {
 
 // autograd methods on torch._C
 static PyMethodDef methods[] = { // NOLINT
-  {"set_grad_enabled", (PyCFunction)set_grad_enabled, METH_O, nullptr},
+  {"_set_grad_enabled", (PyCFunction)set_grad_enabled, METH_O, nullptr},
   {"is_grad_enabled", (PyCFunction)is_grad_enabled, METH_NOARGS, nullptr},
   {"set_autocast_enabled", (PyCFunction)set_autocast_enabled, METH_O, nullptr},
   {"is_autocast_enabled", (PyCFunction)is_autocast_enabled, METH_NOARGS, nullptr},