grab parameter from node

Summary:
Helper functions to:
1. Determine if a placeholder node is a parameter
2. Graph the parameter from a node if it exists

Reviewed By: cccclai

Differential Revision: D47924708

fbshipit-source-id: 38671d87d9826956a7995f8dcbb8ea621a4645e8

Author: mcr229

exir/program/__init__.py

@@ -11,6 +11,8 @@
     edge_to_executorch_passes,
     ExecutorchProgram,
     ExirExportedProgram,
+    get_param,
+    is_param,
     multi_method_program_to_executorch,
     MultiMethodExecutorchProgram,
     MultiMethodExirExportedProgram,
@@ -21,6 +23,8 @@
     "ExecutorchProgram",
     "_to_edge",
     "edge_to_executorch_passes",
+    "is_param",
+    "get_param",
     "MultiMethodExirExportedProgram",
     "MultiMethodExecutorchProgram",
     "multi_method_program_to_executorch",

exir/program/_program.py

@@ -203,6 +203,28 @@ def edge_to_executorch_passes(config: ExecutorchBackendConfig) -> List[PassType]
     return passes
 
 
+def is_param(edge_program: ExportedProgram, node: torch.fx.Node) -> bool:
+    """
+    Checks if the given node is a parameter within the edge_program
+    """
+    return node.name in edge_program.graph_signature.inputs_to_parameters
+
+
+def get_param(
+    edge_program: ExportedProgram,
+    node: torch.fx.Node,
+) -> Optional[torch.nn.Parameter]:
+    """
+    Returns the parameter associated with the given node in the edge program.
+    Returns None if the node is not a parameter within the edge_program
+    """
+    if is_param(edge_program, node):
+        parameter_name = edge_program.graph_signature.inputs_to_parameters[node.name]
+        return edge_program.state_dict[parameter_name]
+
+    return None
+
+
 ######## MULTI METHOD STUFF BELOW HERE. TO BE MERGED INTO ExirExportedProgram and ExecutorchProgram AND THEN DELETED ##########