Allow for HOP to be in the etreord graph

Differential Revision: D80118856

Pull Request resolved: pytorch#13385

Author: kimishpatel

devtools/debug_format/et_schema.py

@@ -29,6 +29,11 @@
     OperatorNode,
     ValueNode,
 )
+
+from torch._higher_order_ops.auto_functionalize import (
+    auto_functionalized,
+    auto_functionalized_v2,
+)
 from torch._subclasses import FakeTensor
 
 
@@ -121,6 +126,12 @@ def _parse_args(  # noqa: C901
             # pyre-ignore
             named_args = node.target._schema.arguments
 
+        if node.op == "call_function" and (
+            node.target == auto_functionalized or node.target == auto_functionalized_v2
+        ):
+            # for functioanlized HOPs, args for the corresponding functional op are stored in kwargs
+            args = tuple(kwargs.values())
+
         for index, arg in enumerate(args):
             if isinstance(arg, torch.fx.node.Node):
                 if arg.target == exir.memory.alloc: