[Reland] Change 'python mode' to 'torch dispatch mode'

Changes Python Mode name to Torch Dispatch Mode because there is now a Torch Function Mode, so Torch Dispatch Mode and Torch Function Mode are consistent with each other
Pull Request resolved: pytorch#76562
Approved by: https://github.com/zou3519, https://github.com/albanD

Author: samdow

aten/src/ATen/ThreadLocalState.cpp

@@ -19,7 +19,7 @@ ThreadLocalState::ThreadLocalState()
 
   saved_tensors_default_hooks_ = at::SavedTensorDefaultHooks::get_stack();
 
-  python_mode_state_ = at::impl::PythonModeTLS::get_state();
+  torch_dispatch_mode_state_ = at::impl::TorchDispatchModeTLS::get_state();
 }
 
 void ThreadLocalState::set_grad_mode(bool enabled) {
@@ -33,7 +33,7 @@ void ThreadLocalState::setThreadLocalState(
   // restore the dispatch key set TLS at the same time.
   c10::AutogradState::set_tls_state(state.autograd_tls_);
 
-  at::impl::PythonModeTLS::set_state(state.python_mode_state_);
+  at::impl::TorchDispatchModeTLS::set_state(state.torch_dispatch_mode_state_);
 
   at::impl::PythonTorchFunctionTLS::set_state(state.python_torch_function_state_);
 

aten/src/ATen/ThreadLocalState.h

@@ -9,7 +9,7 @@
 
 #include <ATen/record_function.h>
 #include <ATen/FuncTorchTLS.h>
-#include <ATen/core/PythonModeTLS.h>
+#include <ATen/core/TorchDispatchModeTLS.h>
 #include <ATen/PythonTorchFunctionTLS.h>
 
 namespace at {
@@ -54,8 +54,8 @@ class TORCH_API ThreadLocalState {
   // TLS for AutogradModes
   AutogradState autograd_tls_;
 
-  // TLS for enable_python_mode (__torch_dispatch__)
-  std::shared_ptr<SafePyObject> python_mode_state_;
+  // TLS for enable_torch_dispatch_mode
+  std::shared_ptr<SafePyObject> torch_dispatch_mode_state_;
 
   // TLS for __torch_function__ (mode and disable_torch_function)
   at::impl::PythonTorchFunctionTLS python_torch_function_state_;

aten/src/ATen/core/PythonFallbackKernel.cpp

@@ -1,4 +1,4 @@
-#include <ATen/core/PythonModeTLS.h>
+#include <ATen/core/TorchDispatchModeTLS.h>
 #include <ATen/core/PythonFallbackKernel.h>
 #include <c10/core/SafePyObject.h>
 
@@ -50,10 +50,10 @@ void pythonFallback(const c10::OperatorHandle& op, torch::jit::Stack* stack) {
   c10::impl::ExcludeDispatchKeyGuard guard(after_Python_keyset);
 
 
-  // If Python Mode is active, use its PyInterpreter for dispatch
-  const auto& maybe_python_mode_state = at::impl::PythonModeTLS::get_state();
-  if (maybe_python_mode_state) {
-    maybe_python_mode_state->pyinterpreter()->dispatch(op, stack, maybe_python_mode_state);
+  // If Torch Dispatch Mode is active, use its PyInterpreter for dispatch
+  const auto& maybe_torch_dispatch_mode_state = at::impl::TorchDispatchModeTLS::get_state();
+  if (maybe_torch_dispatch_mode_state) {
+    maybe_torch_dispatch_mode_state->pyinterpreter()->dispatch(op, stack, maybe_torch_dispatch_mode_state);
     return;
   }
 

aten/src/ATen/core/PythonModeTLS.cpp renamed to aten/src/ATen/core/TorchDispatchModeTLS.cpp

@@ -1,26 +1,26 @@
-#include <ATen/core/PythonModeTLS.h>
+#include <ATen/core/TorchDispatchModeTLS.h>
 #include <c10/core/SafePyObject.h>
 
 namespace at { namespace impl {
 
-thread_local std::shared_ptr<SafePyObject> pythonModeState;
+thread_local std::shared_ptr<SafePyObject> torchDispatchModeState;
 
-void PythonModeTLS::set_state(std::shared_ptr<SafePyObject> state) {
+void TorchDispatchModeTLS::set_state(std::shared_ptr<SafePyObject> state) {
   if (state) {
     c10::impl::tls_set_dispatch_key_included(DispatchKey::Python, true);
     c10::impl::tls_set_dispatch_key_included(DispatchKey::PythonTLSSnapshot, true);
   } else {
-    PythonModeTLS::reset_state();
+    TorchDispatchModeTLS::reset_state();
   }
-  pythonModeState = std::move(state);
+  torchDispatchModeState = std::move(state);
 }
 
-const std::shared_ptr<SafePyObject>& PythonModeTLS::get_state() {
-  return pythonModeState;
+const std::shared_ptr<SafePyObject>& TorchDispatchModeTLS::get_state() {
+  return torchDispatchModeState;
 }
 
-void PythonModeTLS::reset_state() {
-  pythonModeState.reset();
+void TorchDispatchModeTLS::reset_state() {
+  torchDispatchModeState.reset();
   c10::impl::tls_set_dispatch_key_included(DispatchKey::Python, false);
   c10::impl::tls_set_dispatch_key_included(DispatchKey::PythonTLSSnapshot, false);
 }

aten/src/ATen/core/PythonModeTLS.h renamed to aten/src/ATen/core/TorchDispatchModeTLS.h

@@ -7,7 +7,7 @@
 namespace at {
 namespace impl {
 
-struct TORCH_API PythonModeTLS {
+struct TORCH_API TorchDispatchModeTLS {
   static void set_state(std::shared_ptr<SafePyObject> state);
   static const std::shared_ptr<SafePyObject>& get_state();
   static void reset_state();

test/test_decomp.py

@@ -3,7 +3,7 @@
 from collections import defaultdict
 from torch import Tensor
 import torch.autograd
-from torch.utils._python_dispatch import enable_python_mode
+from torch.utils._python_dispatch import enable_torch_dispatch_mode
 from torch._decomp import decomposition_table
 
 from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
@@ -450,7 +450,7 @@ def check_decomposed(aten_name):
                 # explicit clearing is necessary as I will create a fresh mode
                 # for each region
                 decomposed.clear()
-                with enable_python_mode(DecompCrossRefMode):
+                with enable_torch_dispatch_mode(DecompCrossRefMode):
                     decomp_out, decomp_vjp_fn = ref_vjp_no_create(fn, *primals)
                 if aten_name in decomposition_names:
                     check_decomposed(aten_name)
@@ -459,7 +459,7 @@ def check_decomposed(aten_name):
                     cotangents = tree_map(lambda x: torch.randn_like(x), decomp_out)
 
                     decomposed.clear()
-                    with enable_python_mode(DecompCrossRefMode):
+                    with enable_torch_dispatch_mode(DecompCrossRefMode):
                         decomp_vjp_fn(cotangents)
                     if not run_all:
                         check_decomposed(op.aten_backward_name)
@@ -468,7 +468,7 @@ def check_decomposed(aten_name):
                 args = [sample_input.input] + list(sample_input.args)
                 kwargs = sample_input.kwargs
                 decomposed.clear()
-                with enable_python_mode(DecompCrossRefMode):
+                with enable_torch_dispatch_mode(DecompCrossRefMode):
                     func(*args, **kwargs)
                 if not run_all:
                     check_decomposed(aten_name)

test/test_python_dispatch.py

@@ -7,7 +7,7 @@
 from torch.testing._internal.logging_tensor import LoggingTensor, LoggingTensorReentrant, LoggingTensorMode, \
     log_input, capture_logs, no_dispatch
 from torch.utils._pytree import tree_map
-from torch.utils._python_dispatch import enable_python_mode
+from torch.utils._python_dispatch import enable_torch_dispatch_mode
 
 import logging
 
@@ -447,28 +447,28 @@ def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
         res = x.index_put_(idxs, v)
         self.assertEqual(called_funcs, [torch.ops.aten.index_put_.default])
 
-    def test_enable_python_mode_error(self) -> None:
+    def test_enable_torch_dispatch_mode_error(self) -> None:
         with self.assertRaisesRegex(ValueError, "__torch_dispatch__"):
-            with enable_python_mode(torch.Tensor):
+            with enable_torch_dispatch_mode(torch.Tensor):
                 pass
         z = LoggingTensor(torch.empty([]))
         with self.assertRaisesRegex(ValueError, "must be the type"):
-            with enable_python_mode(z):
+            with enable_torch_dispatch_mode(z):
                 pass
 
-    def test_enable_python_mode_basic(self) -> None:
-        with enable_python_mode(LoggingTensorMode):
+    def test_enable_torch_dispatch_mode_basic(self) -> None:
+        with enable_torch_dispatch_mode(LoggingTensorMode):
             z = torch.empty([])
             self.assertTrue(isinstance(z, LoggingTensorMode))
 
-    def test_enable_python_mode_unrelated_tensors(self) -> None:
+    def test_enable_torch_dispatch_mode_unrelated_tensors(self) -> None:
         x = torch.randn([])
         y = torch.randn([])
-        with enable_python_mode(LoggingTensorMode):
+        with enable_torch_dispatch_mode(LoggingTensorMode):
             z = x + y
             self.assertTrue(isinstance(z, LoggingTensorMode))
 
-    def test_enable_python_mode_subclass_priority(self) -> None:
+    def test_enable_torch_dispatch_mode_subclass_priority(self) -> None:
         class ErrorA(RuntimeError):
             pass
 
@@ -500,30 +500,30 @@ def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
 
         # B has precedence over A due to the subclass relationship
         with self.assertRaises(ErrorB):
-            with enable_python_mode(A):
+            with enable_torch_dispatch_mode(A):
                 b + b
         with self.assertRaises(ErrorB):
-            with enable_python_mode(B):
+            with enable_torch_dispatch_mode(B):
                 a + a
         with self.assertRaises(ErrorB):
-            with enable_python_mode(B):
+            with enable_torch_dispatch_mode(B):
                 a + b
 
-    def test_enable_python_mode_respects_no_dispatch(self) -> None:
-        with enable_python_mode(LoggingTensorMode):
+    def test_enable_torch_dispatch_mode_respects_no_dispatch(self) -> None:
+        with enable_torch_dispatch_mode(LoggingTensorMode):
             z = torch.ones([2, 3])
             self.assertTrue(isinstance(z, LoggingTensorMode))
             with no_dispatch():
                 expected = torch.ones([2, 3])
                 self.assertEqual(z.elem, expected)
 
-    def test_nested_enable_python_mode(self) -> None:
+    def test_nested_enable_torch_dispatch_mode(self) -> None:
         with self.assertRaisesRegex(RuntimeError, "has already been set"):
-            with enable_python_mode(LoggingTensorMode):
-                with enable_python_mode(LoggingTensorMode):
+            with enable_torch_dispatch_mode(LoggingTensorMode):
+                with enable_torch_dispatch_mode(LoggingTensorMode):
                     pass
 
-    def test_tolist_numpy_with_python_mode(self) -> None:
+    def test_tolist_numpy_with_torch_dispatch_mode(self) -> None:
         x = LoggingTensor(torch.tensor([2.0, 3.0]))
         with self.assertRaisesRegex(RuntimeError, "is not supported for tensor subclasses."):
             x.tolist()
@@ -532,7 +532,7 @@ def test_tolist_numpy_with_python_mode(self) -> None:
         with self.assertRaises(AssertionError):
             self.assertEqual(x, None)
 
-    def test_enable_python_mode_subclass_autograd_device_check(self) -> None:
+    def test_enable_torch_dispatch_mode_subclass_autograd_device_check(self) -> None:
         class NonWrapperSubclass(torch.Tensor):
             elem: torch.Tensor
 
@@ -554,7 +554,7 @@ def unwrap(e):
                 def wrap(e):
                     return NonWrapperSubclass(e) if isinstance(e, torch.Tensor) else e
 
-                # no_dispatch is only needed if you use enable_python_mode.
+                # no_dispatch is only needed if you use enable_torch_dispatch_mode.
                 # It prevents infinite recursion.
                 with no_dispatch():
                     rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs)))
@@ -591,7 +591,7 @@ def unwrap(e):
                 def wrap(e):
                     return SubclassWithNone(e) if isinstance(e, torch.Tensor) else e
 
-                # no_dispatch is only needed if you use enable_python_mode.
+                # no_dispatch is only needed if you use enable_torch_dispatch_mode.
                 # It prevents infinite recursion.
                 with no_dispatch():
                     rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs)))
@@ -616,7 +616,7 @@ def wrap(e):
             out.backward()
 
     def test_storage_can_be_converted_to_python_object(self):
-        with enable_python_mode(LoggingTensorMode):
+        with enable_torch_dispatch_mode(LoggingTensorMode):
             s = torch.Storage()
             z = LoggingTensorMode(torch.empty([]))
             z.set_(s)

tools/build_variables.bzl

@@ -898,7 +898,7 @@ libtorch_python_core_sources = [
     "torch/csrc/autograd/profiler_python.cpp",
     "torch/csrc/autograd/python_anomaly_mode.cpp",
     "torch/csrc/autograd/python_saved_variable_hooks.cpp",
-    "torch/csrc/autograd/python_mode.cpp",
+    "torch/csrc/autograd/torch_dispatch_mode.cpp",
     "torch/csrc/autograd/python_cpp_function.cpp",
     "torch/csrc/autograd/python_engine.cpp",
     "torch/csrc/autograd/python_function.cpp",
@@ -1081,7 +1081,7 @@ aten_cpu_source_non_codegen_list = [
     "aten/src/ATen/core/op_registration/infer_schema.cpp",
     "aten/src/ATen/core/op_registration/op_registration.cpp",
     "aten/src/ATen/core/operator_name.cpp",
-    "aten/src/ATen/core/PythonModeTLS.cpp",
+    "aten/src/ATen/core/TorchDispatchModeTLS.cpp",
     "aten/src/ATen/core/register_symbols.cpp",
     "aten/src/ATen/core/class_type.cpp",
     "aten/src/ATen/core/type.cpp",

torch/_C/__init__.pyi.in

@@ -747,8 +747,8 @@ def __set_forward_AD_enabled(enabled: _bool) -> None: ...
 def __is_forward_AD_enabled() -> _bool: ...
 def _register_default_hooks(pack_hook: Callable, unpack_hook: Callable) -> None: ...
 def _reset_default_hooks() -> None: ...
-def _enter_python_mode(cls: Type) -> None: ...
-def _exit_python_mode() -> None: ...
+def _enter_torch_dispatch_mode(cls: Type) -> None: ...
+def _exit_torch_dispatch_mode() -> None: ...
 
 class _InferenceMode(object):
     def __init__(self, mode: _bool) -> None: ...

torch/csrc/autograd/init.cpp

@@ -19,7 +19,7 @@
 #include <torch/csrc/autograd/python_saved_variable_hooks.h>
 #include <torch/csrc/autograd/utils/wrap_outputs.h>
 #include <torch/csrc/autograd/utils/python_arg_parsing.h>
-#include <torch/csrc/autograd/python_mode.h>
+#include <torch/csrc/autograd/torch_dispatch_mode.h>
 #include <torch/csrc/autograd/python_variable.h>
 #include <torch/csrc/autograd/record_function_ops.h>
 #include <torch/csrc/utils/pycfunction_helpers.h>
@@ -603,16 +603,16 @@ static PyObject * python_exit_dual_level(PyObject* _unused, PyObject* args, PyOb
   END_HANDLE_TH_ERRORS
 }
 
-static PyObject * enter_python_mode(PyObject* _unused, PyObject* arg) {
+static PyObject * enter_torch_dispatch_mode(PyObject* _unused, PyObject* arg) {
   HANDLE_TH_ERRORS
-  PythonMode::enter(arg);
+  TorchDispatchMode::enter(arg);
   Py_RETURN_NONE;
   END_HANDLE_TH_ERRORS
 }
 
-static PyObject * exit_python_mode(PyObject* _unused, PyObject* arg) {
+static PyObject * exit_torch_dispatch_mode(PyObject* _unused, PyObject* arg) {
   HANDLE_TH_ERRORS
-  PythonMode::exit();
+  TorchDispatchMode::exit();
   Py_RETURN_NONE;
   END_HANDLE_TH_ERRORS
 }
@@ -664,8 +664,8 @@ static PyMethodDef methods[] = { // NOLINT
   {"is_anomaly_enabled", is_anomaly_mode_enabled, METH_NOARGS, nullptr},
   {"_enter_dual_level", python_enter_dual_level, METH_NOARGS, nullptr},
   {"_exit_dual_level", castPyCFunctionWithKeywords(python_exit_dual_level), METH_VARARGS | METH_KEYWORDS, nullptr},
-  {"_enter_python_mode", enter_python_mode, METH_O, nullptr},
-  {"_exit_python_mode", exit_python_mode, METH_NOARGS, nullptr},
+  {"_enter_torch_dispatch_mode", enter_torch_dispatch_mode, METH_O, nullptr},
+  {"_exit_torch_dispatch_mode", exit_torch_dispatch_mode, METH_NOARGS, nullptr},
   {"_set_torch_function_mode", set_torch_function_mode, METH_O, nullptr},
   {"_get_torch_function_mode", get_torch_function_mode, METH_NOARGS, nullptr},
   {nullptr, nullptr, 0, nullptr}

torch/csrc/autograd/python_mode.cpp

@@ -33,7 +33,7 @@
 
 #include <torch/library.h>
 #include <torch/csrc/jit/python/pybind_utils.h>
-#include <torch/csrc/autograd/python_mode.h>
+#include <torch/csrc/autograd/torch_dispatch_mode.h>
 
 
 #include <ATen/ATen.h>
@@ -1739,9 +1739,9 @@ bool isPythonTensor(const Tensor& tensor) {
 }
 
 // NOTE [dispatch_fn's type argument]
-// `type` is nullable and represents the PythonMode going on.
-// Right now we only support a single PythonMode, but in the future we could
-// change this to a stack of PythonModes.
+// `type` is nullable and represents the TorchDispatchMode going on.
+// Right now we only support a single TorchDispatchMode, but in the future we could
+// change this to a stack of TorchDispatchModes.
 //
 // If `type` isn't null, then we consider the type for dispatch by prepending
 // it to the overloaded_args list. `handle_torch_funciton_no_python_arg_parser`