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

Author: xilzy

graph_net/torch/backend/unstable_to_stable_backend.py

@@ -161,7 +161,31 @@ def _impl_unstable_to_stable_special_logit(self, gm):
 
     # replace this line with modification code for task 123 (torch._C._nn.pad)
 
-    # 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

@@ -32,7 +32,7 @@ def serialize_graph_module_to_str(gm: torch.fx.GraphModule) -> str:
         # replace this line with modification code for task 121 (torch._C._set_grad_enabled)
         # replace this line with modification code for task 122 (torch._C._log_api_usage_once)
         # replace this line with modification code for task 123 (torch._C._nn.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)
         # replace this line with modification code for task 127 (torch._C._nn.linear)
     ]