(node: torch.fx.Node)
| 462 | return list(args) |
| 463 | |
| 464 | def extract_input(node: torch.fx.Node) -> Optional[FakeTensor]: |
| 465 | if "val" in node.meta: |
| 466 | fake = node.meta["val"] |
| 467 | if hasattr(fake, "constant") and fake.constant is not None: |
| 468 | return fake.constant |
| 469 | return fake |
| 470 | elif tensor_meta := node.meta.get("tensor_meta"): |
| 471 | assert self.fake_tensor_mode is not None |
| 472 | return FakeTensor( |
| 473 | self.fake_tensor_mode, |
| 474 | torch.empty( |
| 475 | tensor_meta.shape, |
| 476 | dtype=tensor_meta.dtype, |
| 477 | device="meta", |
| 478 | requires_grad=tensor_meta.requires_grad, |
| 479 | memory_format=tensor_meta.memory_format, |
| 480 | ), |
| 481 | torch.device("cpu"), |
| 482 | ) |
| 483 | elif len(node.users) == 0: |
| 484 | return None |
| 485 | raise ExportPassBaseError( |
| 486 | f"Cannot construct an input for graph module: {graph_module}.", |
| 487 | ) |
| 488 | |
| 489 | return [ |
| 490 | extract_input(node) |
nothing calls this directly
no test coverage detected