[HigherOrderOp] Remove _deprecated_global_ns from some ops (#104105)

zou3519 · pytorchmergebot · commit 280df5dc2e0a · 2023-06-28T00:03:29.000Z
The remaining ops after this PR are: - cond - map - anything that is out of tree. These are a bit more difficult to remove. Test Plan: - existing tests Pull Request resolved: #104105 Approved by: https://github.com/ydwu4
diff --git a/torch/_higher_order_ops/wrap.py b/torch/_higher_order_ops/wrap.py
@@ -6,7 +6,7 @@
 # Used for testing the HigherOrderOperator mechanism
 class Wrap(HigherOrderOperator):
     def __init__(self):
-        super().__init__("wrap", _deprecated_global_ns=True)
+        super().__init__("wrap")
 
     def __call__(self, func, *args):
         # Dynamo already traces the body of HigherOrderOp beforehand when it
@@ -41,7 +41,7 @@ class WrapActivationCheckpoint(HigherOrderOperator):
     partitioners. See TagActivationCheckpoint for more information.
     """
     def __init__(self):
-        super().__init__("wrap_activation_checkpoint", _deprecated_global_ns=True)
+        super().__init__("wrap_activation_checkpoint")
 
     def __call__(self, function, *args, **kwargs):
         # use_reentrant is set to False because this op is going to be traced.
@@ -72,7 +72,7 @@ class TagActivationCheckpoint(HigherOrderOperator):
     """
 
     def __init__(self):
-        super().__init__("tag_activation_checkpoint", _deprecated_global_ns=True)
+        super().__init__("tag_activation_checkpoint")
 
     def tag_nodes(self, gmod):
         # TODO - This needs major investigation. Currently, we are tagging all
diff --git a/torch/_prims/rng_prims.py b/torch/_prims/rng_prims.py
@@ -164,9 +164,7 @@ def get_device(args, kwargs):
 
 
 def register_run_and_save_rng_state_op():
-    run_and_save_rng_state = HigherOrderOperator(
-        "run_and_save_rng_state", _deprecated_global_ns=True
-    )
+    run_and_save_rng_state = HigherOrderOperator("run_and_save_rng_state")
 
     run_and_save_rng_state.fallthrough(DispatchKey.ADInplaceOrView)
     run_and_save_rng_state.fallthrough(DispatchKey.PythonDispatcher)  # type: ignore[attr-defined]
@@ -220,9 +218,7 @@ def impl_proxy_dispatch_mode(op, *args, **kwargs):
 
 
 def register_run_with_rng_state_op():
-    run_with_rng_state = HigherOrderOperator(
-        "run_with_rng_state", _deprecated_global_ns=True
-    )
+    run_with_rng_state = HigherOrderOperator("run_with_rng_state")
 
     run_with_rng_state.fallthrough(DispatchKey.ADInplaceOrView)
     run_with_rng_state.fallthrough(DispatchKey.PythonTLSSnapshot)  # type: ignore[attr-defined]
