| 106 | |
| 107 | |
| 108 | def get_param_tensor( |
| 109 | exp_prog: ExportedProgram, node: torch.fx.Node |
| 110 | ) -> Optional[torch.Tensor]: |
| 111 | if node is None: |
| 112 | return None |
| 113 | elif is_param(exp_prog, node): |
| 114 | return get_param(exp_prog, node) |
| 115 | elif is_buffer(exp_prog, node): |
| 116 | return get_buffer(exp_prog, node) |
| 117 | elif is_lifted_tensor_constant(exp_prog, node): |
| 118 | return get_lifted_tensor_constant(exp_prog, node) |
| 119 | elif is_get_attr_node(node): |
| 120 | # This is a hack to support both lifted and unlifted graph |
| 121 | try: |
| 122 | return getattr(node.graph.owning_module, node.target) |
| 123 | except AttributeError: |
| 124 | return getattr(exp_prog.graph_module, node.target) |
| 125 | raise RuntimeError(f"unsupported param type, {node.op}.") |
| 126 | |
| 127 | |
| 128 | def get_tensor_name(exp_prog: ExportedProgram, node: torch.fx.Node) -> str: |