Creates and returns a constant placeholder node, meaning that it is of type parameter, buffer, or lifted constant tensor. graph.inserting_before/after() should be used before the call to decide where to insert the node, at an insertion point before the first input node.
(
exp_program: ExportedProgram,
graph: torch.fx.Graph,
name: str,
kind: InputKind,
data: torch.Tensor,
persistent_buffer: Optional[bool] = None,
)
| 96 | |
| 97 | |
| 98 | def create_constant_placeholder( |
| 99 | exp_program: ExportedProgram, |
| 100 | graph: torch.fx.Graph, |
| 101 | name: str, |
| 102 | kind: InputKind, |
| 103 | data: torch.Tensor, |
| 104 | persistent_buffer: Optional[bool] = None, |
| 105 | ) -> torch.fx.Node: |
| 106 | """ |
| 107 | Creates and returns a constant placeholder node, meaning that it is of type parameter, buffer, |
| 108 | or lifted constant tensor. graph.inserting_before/after() should be used before the call to |
| 109 | decide where to insert the node, at an insertion point before the first input node. |
| 110 | """ |
| 111 | |
| 112 | target = name |
| 113 | |
| 114 | # If a placeholder with this target already exists, return it to avoid |
| 115 | # duplicate parameter names in the generated function signature which would |
| 116 | # cause a SyntaxError on recompile. This can happen when multiple pattern |
| 117 | # replacements independently create placeholders for a shared weight. |
| 118 | if name in exp_program.state_dict or name in exp_program.constants: |
| 119 | for n in graph.nodes: |
| 120 | if n.op == "placeholder" and n.target == name: |
| 121 | return n |
| 122 | |
| 123 | # Add data to state_dict/ constants |
| 124 | match kind: |
| 125 | case InputKind.PARAMETER: |
| 126 | exp_program.state_dict[target] = torch.nn.Parameter( |
| 127 | data, requires_grad=False |
| 128 | ) |
| 129 | case InputKind.BUFFER: |
| 130 | if persistent_buffer is None: |
| 131 | raise RuntimeError( |
| 132 | "Must set persistent_buffer when creating a new buffer." |
| 133 | ) |
| 134 | elif persistent_buffer: |
| 135 | exp_program.state_dict[target] = data |
| 136 | else: |
| 137 | exp_program.constants[target] = data |
| 138 | case InputKind.CONSTANT_TENSOR: |
| 139 | exp_program.constants[target] = data |
| 140 | case _: |
| 141 | raise RuntimeError("Can only create constant input nodes.") |
| 142 | |
| 143 | fake_tensor = _get_fake_tensor_mode(graph, data) |
| 144 | |
| 145 | # Create node |
| 146 | node = graph.create_node(op="placeholder", name=name, target=name) |
| 147 | node.meta["val"] = fake_tensor |
| 148 | |
| 149 | # Add tensor to graph_signature in the same order as nodes in the graph |
| 150 | node_names = [n.name for n in graph.nodes if n.op == "placeholder"] |
| 151 | node_index = node_names.index(name) |
| 152 | |
| 153 | input_specs = exp_program.graph_signature.input_specs |
| 154 | user_input_indices = [ |
| 155 | i for i, spec in enumerate(input_specs) if spec.kind == InputKind.USER_INPUT |