[HigherOrderOp] Remove _deprecated_global_ns from cond (#104380)

ydwu4 · pytorchmergebot · commit 6a3d5f198632 · 2023-07-07T17:13:09.000Z
Remove _deprecated_global_ns from cond following #104105. We change the module attribute of HigherOrderOperator instances in the constructor from torch.ops to torch.ops.higher_order when self.namespace is "higher_order". For subclasses (e.g. customized higher order operator), we leave their \_\_module\_\_ unchanged. Will import this PR to fix internal tests. Pull Request resolved: #104380 Approved by: https://github.com/zhxchen17, https://github.com/zou3519
diff --git a/functorch/experimental/_cond.py b/functorch/experimental/_cond.py
@@ -32,7 +32,7 @@ class UnsupportedAliasMutationException(RuntimeError):
 We're going to define a `cond` operation.
 In order to do this, we need implementations for each of the dispatch keys.
 """
-cond = HigherOrderOperator("cond", _deprecated_global_ns=True)
+cond = HigherOrderOperator("cond")
 
 def trace_cond(proxy_mode, func_overload, pred, true_fn, false_fn, operands):
     assert isinstance(operands, (list, tuple)), "Cond operands must be a list or tuple of tensors"
diff --git a/test/dynamo/test_export.py b/test/dynamo/test_export.py
@@ -3209,7 +3209,7 @@ def forward(self, pred, x):
     ones_3 = torch.ones(6, 4)
     cond_true_0 = self.cond_true_0
     cond_false_0 = self.cond_false_0
-    cond = torch.ops.cond(arg0, cond_true_0, cond_false_0, [arg1, ones, ones_1, ones_3, ones, ones_1, ones_2]);  arg0 = cond_true_0 = cond_false_0 = arg1 = ones = ones_1 = ones_3 = ones_2 = None
+    cond = torch.ops.higher_order.cond(arg0, cond_true_0, cond_false_0, [arg1, ones, ones_1, ones_3, ones, ones_1, ones_2]);  arg0 = cond_true_0 = cond_false_0 = arg1 = ones = ones_1 = ones_3 = ones_2 = None
     return pytree.tree_unflatten([cond], self._out_spec)""",  # noqa: B950,E122
         )
 
diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py
@@ -613,11 +613,11 @@ def f(x, pred, pred2):
         def forward(self, x_1, pred_1, pred2_1):
             true_graph_0 = self.true_graph_0
             false_graph_0 = self.false_graph_0
-            conditional = torch.ops.cond(pred_1, true_graph_0, false_graph_0, [x_1]);
+            conditional = torch.ops.higher_order.cond(pred_1, true_graph_0, false_graph_0, [x_1]);
             pred_1 = true_graph_0 = false_graph_0 = None
             true_graph_1 = self.true_graph_1
             false_graph_1 = self.false_graph_1
-            conditional_1 = torch.ops.cond(pred2_1, true_graph_1, false_graph_1, [x_1, x_1]);
+            conditional_1 = torch.ops.higher_order.cond(pred2_1, true_graph_1, false_graph_1, [x_1, x_1]);
             pred2_1 = true_graph_1 = false_graph_1 = x_1 = None
             add = torch.ops.aten.add.Tensor(conditional, conditional_1);  conditional = conditional_1 = None
             return add
@@ -780,11 +780,11 @@ def f(x, pred, pred2):
         def forward(self, x_1, pred_1, pred2_1):
             true_graph_0 = self.true_graph_0
             false_graph_0 = self.false_graph_0
-            conditional = torch.ops.cond(pred_1, true_graph_0, false_graph_0, [x_1]);
+            conditional = torch.ops.higher_order.cond(pred_1, true_graph_0, false_graph_0, [x_1]);
             pred_1 = true_graph_0 = false_graph_0 = None
             true_graph_1 = self.true_graph_1
             false_graph_1 = self.false_graph_1
-            conditional_1 = torch.ops.cond(pred2_1, true_graph_1, false_graph_1, [x_1, x_1]);
+            conditional_1 = torch.ops.higher_order.cond(pred2_1, true_graph_1, false_graph_1, [x_1, x_1]);
             pred2_1 = true_graph_1 = false_graph_1 = x_1 = None
             add = torch.ops.aten.add.Tensor(conditional, conditional_1);  conditional = conditional_1 = None
             return add
diff --git a/torch/_ops.py b/torch/_ops.py
@@ -214,6 +214,14 @@ def __init__(self, name, *, _deprecated_global_ns=False):
         else:
             _higher_order_ops[name] = self
             self._ns = "higher_order"
+
+        # For a normal HigherOrderOperator instance, we will change its __module__ from torch._ops to
+        # torch._ops.higher_order.
+        # For an instance of subclass of HigherOrderOperator (e.g. customized higher order op),
+        # the __module__ attribute will be kept unchanged.
+        if self.__class__ is HigherOrderOperator:
+            self_name_space = "." + self.namespace if self.namespace else ""
+            self.__module__ = self.__module__ + self_name_space
         self.non_fallthrough_keys = torch._C._dispatch_keyset_full()
 
     @property
@@ -768,7 +776,14 @@ def __init__(self, name, ops):
         self._ops = ops
 
     def __getattr__(self, name):
-        return self._ops[name]
+        # Following _OpNamespace.__getattr__, we cache the op on the _PyOpNamespace object.
+        op = self._ops.get(name, None)
+        if op is None:
+            raise AttributeError(
+                f"'_PyOpNamespace' '{self.name}' object has no attribute '{name}'"
+            )
+        setattr(self, name, op)
+        return op
 
 
 class _Ops(types.ModuleType):

Original file line number	Diff line number	Diff line change
`@@ -3209,7 +3209,7 @@ def forward(self, pred, x):`
`3209`	`3209`	`ones_3 = torch.ones(6, 4)`
`3210`	`3210`	`cond_true_0 = self.cond_true_0`
`3211`	`3211`	`cond_false_0 = self.cond_false_0`
`3212`		`- cond = torch.ops.cond(arg0, cond_true_0, cond_false_0, [arg1, ones, ones_1, ones_3, ones, ones_1, ones_2]); arg0 = cond_true_0 = cond_false_0 = arg1 = ones = ones_1 = ones_3 = ones_2 = None`
	`3212`	`+ cond = torch.ops.higher_order.cond(arg0, cond_true_0, cond_false_0, [arg1, ones, ones_1, ones_3, ones, ones_1, ones_2]); arg0 = cond_true_0 = cond_false_0 = arg1 = ones = ones_1 = ones_3 = ones_2 = None`
`3213`	`3213`	`return pytree.tree_unflatten([cond], self._out_spec)""", # noqa: B950,E122`
`3214`	`3214`	`)`
`3215`	`3215`