Deserialize operator args even if operator wasn't found

Summary:
Today when de-serializing a graph if the operator is not found we just return the fake operator without any args or kwargs. This doesn't work for the SDK visualization purpose as we need to render the graph with all the connections.

In order to support this only for EXIR targeted serialization i'm adding support for deserializing operators (their inputs and outputs) without any schema

Reviewed By: angelayi

Differential Revision: D48258339

fbshipit-source-id: 7601c76c6ada2cc01fdfeadbba6814f40c9275a9

Author: tarun292

exir/serde/serialize.py

@@ -364,6 +364,12 @@ def find_operator(module: _DialectNamespace, serialized_target: str) -> str:
 
         return super().deserialize_operator(serialized_target)
 
+    # pyre-ignore
+    def deserialize_inputs_no_schema(self, serialized_node) -> Any:
+        return tuple(
+            self.deserialize_input(input.arg) for input in serialized_node.inputs
+        )
+
     # pyre-ignore
     def deserialize_node(self, serialized_node: schema.Node, target: Callable) -> None:
         if target == "memory.alloc":
@@ -426,6 +432,11 @@ def fake_op(x):
 
             fake_op.__name__ = target
             target = fake_op
+
+            args = self.deserialize_inputs_no_schema(serialized_node)
+            fx_node = self.graph.create_node("call_function", target, args, None, None)
+            self.deserialize_arbitrary_outputs(serialized_node, fx_node)
+
             return
 
         super().deserialize_node(serialized_node, target)
@@ -502,12 +513,20 @@ def deserialize_alloc_spec(serialized_alloc_spec: str) -> memory.AllocSpec:
     def deserialize_arbitrary_outputs(
         self, serialized_node: schema.Node, fx_node: torch.fx.Node
     ) -> None:
+        if len(serialized_node.outputs) == 0:
+            return
         # Single tensor return
-        if (
+        elif (
             len(serialized_node.outputs) == 1
             and serialized_node.outputs[0].type == "as_tensor"
         ):
             return self.sync_fx_node(serialized_node.outputs[0].as_tensor.name, fx_node)
+        elif len(serialized_node.outputs) == 1 and isinstance(
+            serialized_node.outputs[0].value,
+            (schema.SymIntArgument, schema.SymBoolArgument),
+        ):
+            self.sync_fx_node(serialized_node.outputs[0].value.as_name, fx_node)
+            return
 
         self.deserialize_multiple_outputs(serialized_node, fx_node)
 

sdk/edir/et_schema.py

@@ -486,6 +486,8 @@ def _parse_args(
 
         for arg in args:
             if isinstance(arg, torch.fx.node.Node):
+                if arg.target == exir.memory.alloc:
+                    continue
                 arg_name = FXOperatorGraph._get_node_name(arg)
             elif isinstance(arg, (int, float, torch.dtype)):
                 # e.g. The "0" from node.args of squeeze_copy (mm_default, 0)