Convert torch._C._nn.gelu to torch.nn.functional.gelu

Author: xilzy

graph_net/torch/backend/unstable_to_stable_backend.py

@@ -292,7 +292,31 @@ def replace_in_graph(graph_mod):
 
         return gm
 
-    # replace this line with modification code for task 125 (torch._C._nn.gelu)
+    def _impl_unstable_to_stable_gelu(self, gm):
+        """
+        Convert torch._C._nn.gelu to torch.nn.functional.gelu
+        """
+        import torch.nn.functional as F
+
+        def replace_in_graph(graph_mod):
+            for node in graph_mod.graph.nodes:
+                if node.op == "call_function":
+                    if "gelu" in str(node.target) and "torch._C._nn" in str(
+                        node.target
+                    ):
+                        node.target = F.gelu
+            graph_mod.recompile()
+
+        modules = [gm]
+        modules += [
+            m
+            for _, m in gm.named_modules()
+            if isinstance(m, torch.fx.GraphModule) and m is not gm
+        ]
+        for m in modules:
+            replace_in_graph(m)
+
+        return gm
 
     # replace this line with modification code for task 126 (torch._C._nn.scaled_dot_product_attention)
 

graph_net/torch/fx_graph_serialize_util.py

@@ -147,7 +147,7 @@ def serialize_graph_module_to_str(gm: torch.fx.GraphModule) -> str:
         (r"torch\._C\.set_grad_enabled\(", "torch.set_grad_enabled("),
         # replace this line with modification code for task 122 (torch._C._log_api_usage_once)
         (r"torch\._C\._nn\.pad\(", "torch.nn.functional.pad("),
-        # replace this line with modification code for task 125 (torch._C._nn.gelu)
+        (r"torch\._C\._nn\.gelu\(", "torch.nn.functional.gelu("),
         # replace this line with modification code for task 126 (torch._C._nn.scaled_dot_product_attention)
         (r"torch\._C\._nn\.linear\(", "torch.nn.functional.linear("),
     ]