Creates and returns a mutable buffer placeholder node. This is similar to create_constant_placeholder but specifically for creating mutable buffers that can be modified during execution. The difference between this and create_constant_placeholder is that this doesn't expect use
(
exp_program: ExportedProgram,
name: str,
data: torch.Tensor,
)
| 270 | |
| 271 | |
| 272 | def create_mutable_buffer( |
| 273 | exp_program: ExportedProgram, |
| 274 | name: str, |
| 275 | data: torch.Tensor, |
| 276 | ) -> torch.fx.Node: |
| 277 | """ |
| 278 | Creates and returns a mutable buffer placeholder node. This is similar to |
| 279 | create_constant_placeholder but specifically for creating mutable buffers that |
| 280 | can be modified during execution. |
| 281 | |
| 282 | The difference between this and create_constant_placeholder is that this doesn't |
| 283 | expect user to set the correct position for the placeholder node to be inserted, |
| 284 | it finds the correct position automatically. |
| 285 | |
| 286 | It also updates the graph outputs to include the mutable buffer. |
| 287 | |
| 288 | Args: |
| 289 | exp_program: The exported program to modify |
| 290 | name: The name for the new buffer node (should start with "b_" prefix by convention) |
| 291 | data: The initial tensor data for the buffer |
| 292 | |
| 293 | Returns: |
| 294 | The created placeholder node to be used in the graph |
| 295 | """ |
| 296 | # Input validation |
| 297 | if not name or not name.strip(): |
| 298 | raise ValueError("Buffer name cannot be empty") |
| 299 | |
| 300 | if not isinstance(data, torch.Tensor): |
| 301 | raise ValueError("Data must be a torch.Tensor") |
| 302 | |
| 303 | # Extract target name (remove "b_" prefix if present, following export convention) |
| 304 | if name.startswith("b_"): |
| 305 | target = name[2:] |
| 306 | else: |
| 307 | target = name |
| 308 | |
| 309 | # Check if target already exists |
| 310 | if target in exp_program.state_dict: |
| 311 | raise RuntimeError(f"Buffer target '{target}' already exists in state_dict") |
| 312 | |
| 313 | _validate_graph_signature(exp_program) |
| 314 | |
| 315 | persistent_buffer = True |
| 316 | exp_program.state_dict[target] = data |
| 317 | |
| 318 | graph = exp_program.graph_module.graph |
| 319 | |
| 320 | # Create fake tensor using helper function |
| 321 | fake_tensor = _get_fake_tensor_mode(graph, data) |
| 322 | |
| 323 | # Signature ordering is as follows: |
| 324 | # Inputs = [*parameters_buffers_constant_tensors, *flattened_user_inputs] |
| 325 | # ^^^^^^^ |
| 326 | # insert here (at the end of buffers) |
| 327 | # Outputs = [*mutated_inputs, *flattened_user_outputs] |
| 328 | # ^^^^^^^^^^^^^^^ |
| 329 | # insert here (at the end of mutated inputs) |