Update on "[etrecord] Implement generic fallback for GraphModuleSerializer.handle_call_function"

Title says it all!

Implement the case where `node.target` is neither `torch._ops.OpOverload` or `torch._ops.HigherOrderOperator`, instead of throwing an exception.

Differential Revision: [D88216198](https://our.internmc.facebook.com/intern/diff/D88216198/)

[ghstack-poisoned]

Author: ssjia

exir/serde/serialize.py

@@ -450,19 +450,23 @@ def deserialize_node(self, serialized_node: schema.Node, target: Callable) -> No
             fx_node.meta.update(self.deserialize_metadata(serialized_node.metadata))
             return
         elif isinstance(target, str):
-            # Create a dummy fake op if the target does not exist
-            # because we cannot create a call_function node w/o a
-            # callable target
-            log.warning(
-                f"Could not find operator {target}. Returning fake operator."
-            )  # noqa: G004
-
-            # pyre-ignore
-            def fake_op(x):
-                raise NotImplementedError("Fake op is not meant to be run.")
-
-            fake_op.__name__ = target
-            target = fake_op
+            # Special handling for memory ops, which are not EdgeOpOverload but
+            # are still somewhat expected in serialized graphs.
+            if target == "executorch.exir.memory.view":
+                target = exir.memory.view
+            else:
+                # Otherwise, create a dummy fake op if the target does not exist
+                # because we cannot create a call_function node w/o a callable
+                # target
+                log.warning(
+                    f"Could not find operator {target}. Returning fake operator."
+                )  # noqa: G004
+                # pyre-ignore
+                def fake_op(x):
+                    raise NotImplementedError("Fake op is not meant to be run.")
+
+                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)

exir/tests/test_serde.py

@@ -335,3 +335,18 @@ def forward(self, x):
                     node.meta.get("from_node"), node_new.meta.get("from_node")
                 ):
                     self.assertEqual(node_source.to_dict(), node_source_new.to_dict())
+
+    def test_memory_ops(self) -> None:
+        class MemoryOpsModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x, y):
+                x = exir.memory.view(x, (10, 10))
+                return x + y
+
+        inputs = (
+            torch.randn(100),
+            torch.randn(10, 10),
+        )
+        self.check_serde(MemoryOpsModule(), inputs)