feat: convert torch._C._nn.linear to torch.nn.functional.linear

- Implement direct node.target modification for API conversion
- Use serialize_graph_module_to_str for API check
- Add AST-based replacement function (commented) in fx_graph_serialize_util.py
- Simplify conversion logic by removing complex AST code
- Tested with 50 samples: 100% success rate, ES(-6) = 1.013

Author: Dayuxiaoshui

graph_net/torch/backend/unstable_to_stable_backend.py

@@ -167,14 +167,15 @@ def _impl_unstable_to_stable_linear_to_functional_linear(self, gm):
         Returns:
             Modified GraphModule object
         """
+        import torch.nn.functional as F
+
         # Get reference to torch._C._nn.linear for comparison
         try:
             unstable_linear = torch._C._nn.linear
         except AttributeError:
             unstable_linear = None
 
         # Traverse all nodes to find nodes that need to be replaced
-        nodes_to_replace = []
         for node in gm.graph.nodes:
             if node.op == "call_function":
                 target = node.target
@@ -203,108 +204,10 @@ def _impl_unstable_to_stable_linear_to_functional_linear(self, gm):
                     should_replace = True
 
                 if should_replace:
-                    nodes_to_replace.append(node)
-
-        # Since torch._C._nn.linear and torch.nn.functional.linear are the same object,
-        # the code generator cannot distinguish them, so we need to use AST to modify the code string after code generation
-        if nodes_to_replace:
-            # First recompile to generate code
-            gm.recompile()
-
-            # Use AST to modify the generated code, replacing torch._C._nn.linear with torch.nn.functional.linear
-            code_str = gm.code
-
-            # Parse AST
-            tree = ast.parse(code_str)
-
-            class LinearReplacer(ast.NodeTransformer):
-                def visit_Call(self, node):
-                    # Check if it's a torch._C._nn.linear call
-                    # Structure: torch._C._nn.linear(...)
-                    if isinstance(node.func, ast.Attribute):
-                        # node.func.attr should be "linear"
-                        if node.func.attr == "linear":
-                            # node.func.value should be torch._C._nn
-                            if isinstance(node.func.value, ast.Attribute):
-                                # node.func.value.attr should be "_nn"
-                                if node.func.value.attr == "_nn":
-                                    # node.func.value.value should be torch._C
-                                    if isinstance(node.func.value.value, ast.Attribute):
-                                        # node.func.value.value.attr should be "_C"
-                                        if node.func.value.value.attr == "_C":
-                                            # node.func.value.value.value should be torch
-                                            if (
-                                                isinstance(
-                                                    node.func.value.value.value,
-                                                    ast.Name,
-                                                )
-                                                and node.func.value.value.value.id
-                                                == "torch"
-                                            ):
-                                                # Found torch._C._nn.linear, replace with torch.nn.functional.linear
-                                                new_func = ast.Attribute(
-                                                    value=ast.Attribute(
-                                                        value=ast.Attribute(
-                                                            value=ast.Name(
-                                                                id="torch",
-                                                                ctx=ast.Load(),
-                                                            ),
-                                                            attr="nn",
-                                                            ctx=ast.Load(),
-                                                        ),
-                                                        attr="functional",
-                                                        ctx=ast.Load(),
-                                                    ),
-                                                    attr="linear",
-                                                    ctx=ast.Load(),
-                                                )
-                                                node.func = new_func
-                    return self.generic_visit(node)
-
-            transformer = LinearReplacer()
-            modified_tree = transformer.visit(tree)
-            ast.fix_missing_locations(modified_tree)
-
-            # Convert the modified AST back to code string
-            new_code = ast.unparse(modified_tree)
-
-            # Recompile the modified code
-            # Need to import device, inf and other modules that may be used
-            namespace = {
-                "torch": torch,
-            }
-            # Try to import device (if used in code)
-            try:
-                from torch import device
-
-                namespace["device"] = device
-            except ImportError:
-                pass
-            # Try to import inf (if used in code)
-            try:
-                from torch import inf
-
-                namespace["inf"] = inf
-            except ImportError:
-                # If torch doesn't have inf, use math.inf
-                try:
-                    import math
+                    node.target = F.linear
 
-                    namespace["inf"] = math.inf
-                except:
-                    pass
-
-            exec(compile(modified_tree, filename="<ast>", mode="exec"), namespace)
-
-            # Update GraphModule's forward method and code
-            forward_func = namespace.get("forward")
-            if forward_func:
-                import types
-
-                gm.forward = types.MethodType(forward_func, gm)
-
-            # Update _code attribute so that gm.code returns the modified code
-            gm._code = new_code
+        # Recompile the graph
+        gm.recompile()
 
         return gm
 

graph_net/torch/fx_graph_serialize_util.py

@@ -2,6 +2,120 @@
 import torch.fx
 
 
+# def apply_ast_based_linear_replacement(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
+#     """
+#     Apply AST-based replacement of torch._C._nn.linear to torch.nn.functional.linear.
+#
+#     This function uses AST parsing and transformation to replace torch._C._nn.linear
+#     calls with torch.nn.functional.linear in the GraphModule's code.
+#
+#     Note: This function is currently commented out as the replacement is now handled
+#     by simple string replacement in serialize_graph_module_to_str.
+#
+#     Args:
+#         gm: The GraphModule to modify.
+#
+#     Returns:
+#         Modified GraphModule with torch._C._nn.linear replaced by torch.nn.functional.linear.
+#     """
+#     import ast
+#     import torch
+#     import types
+#
+#     # First recompile to generate code
+#     gm.recompile()
+#
+#     # Use AST to modify the generated code, replacing torch._C._nn.linear with torch.nn.functional.linear
+#     code_str = gm.code
+#
+#     # Parse AST
+#     tree = ast.parse(code_str)
+#
+#     class LinearReplacer(ast.NodeTransformer):
+#         def visit_Call(self, node):
+#             # Check if it's a torch._C._nn.linear call
+#             # Structure: torch._C._nn.linear(...)
+#             filtered_nodes = [
+#                 node
+#                 for node in [node]
+#                 if isinstance(node.func, ast.Attribute)
+#                 if node.func.attr == "linear"
+#                 if isinstance(node.func.value, ast.Attribute)
+#                 if node.func.value.attr == "_nn"
+#                 if isinstance(node.func.value.value, ast.Attribute)
+#                 if node.func.value.value.attr == "_C"
+#                 if isinstance(node.func.value.value.value, ast.Name)
+#                 if node.func.value.value.value.id == "torch"
+#             ]
+#             if filtered_nodes:
+#                 # Found torch._C._nn.linear, replace with torch.nn.functional.linear
+#                 new_func = ast.Attribute(
+#                     value=ast.Attribute(
+#                         value=ast.Attribute(
+#                             value=ast.Name(
+#                                 id="torch",
+#                                 ctx=ast.Load(),
+#                             ),
+#                             attr="nn",
+#                             ctx=ast.Load(),
+#                         ),
+#                         attr="functional",
+#                         ctx=ast.Load(),
+#                     ),
+#                     attr="linear",
+#                     ctx=ast.Load(),
+#                 )
+#                 node.func = new_func
+#             return self.generic_visit(node)
+#
+#     transformer = LinearReplacer()
+#     modified_tree = transformer.visit(tree)
+#     ast.fix_missing_locations(modified_tree)
+#
+#     # Convert the modified AST back to code string
+#     new_code = ast.unparse(modified_tree)
+#
+#     # Recompile the modified code
+#     # Need to import device, inf and other modules that may be used
+#     namespace = {
+#         "torch": torch,
+#     }
+#     # Try to import device (if used in code)
+#     try:
+#         from torch import device
+#
+#         namespace["device"] = device
+#     except ImportError:
+#         pass
+#     # Try to import inf (if used in code)
+#     try:
+#         from torch import inf
+#
+#         namespace["inf"] = inf
+#     except ImportError:
+#         # If torch doesn't have inf, use math.inf
+#         try:
+#             import math
+#
+#             namespace["inf"] = math.inf
+#         except:
+#             pass
+#
+#     exec(compile(modified_tree, filename="<ast>", mode="exec"), namespace)
+#
+#     # Update GraphModule's forward method
+#     forward_func = namespace.get("forward")
+#     if forward_func:
+#         gm.forward = types.MethodType(forward_func, gm)
+#
+#     # Use serialize_graph_module_to_str to get the serialized code
+#     # This ensures the code is properly serialized with unstable API replacements
+#     serialized_code = serialize_graph_module_to_str(gm)
+#     gm._code = serialized_code
+#
+#     return gm
+
+
 def serialize_graph_module_to_str(gm: torch.fx.GraphModule) -> str:
     """
     Serialize a GraphModule to a string representation, replacing unstable APIs
@@ -34,7 +148,7 @@ def serialize_graph_module_to_str(gm: torch.fx.GraphModule) -> str:
         # 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)
         # 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)
+        (r"torch\._C\._nn\.linear\(", "torch.nn.functional.linear("),
     ]
     for pattern, repl in replacements:
         code = re.sub(pattern, repl, code)