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

Author: xilzy

graph_net/torch/backend/unstable_to_stable_backend.py

@@ -159,7 +159,29 @@ def _impl_unstable_to_stable_special_logit(self, gm):
 
     # 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)
+    def _impl_unstable_to_stable_pad(self, gm):
+        """
+        Convert torch._C._nn.pad to torch.nn.functional.pad
+        """
+        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 "pad" in str(node.target) and "torch._C._nn" in str(node.target):
+                        node.target = F.pad
+            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 125 (torch._C._nn.gelu)
 

graph_net/torch/fx_graph_serialize_util.py

@@ -31,7 +31,7 @@ def serialize_graph_module_to_str(gm: torch.fx.GraphModule) -> str:
         # replace this line with modification code for task 119 (torch._C._nn.one_hot)
         # 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)
+        (r"torch\._C\._nn\.pad\(", "torch.nn.functional.pad("),
         # replace this line with modification code for task 125 (torch._C._nn.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)