Helper function to create a fake tensor using the fake_mode from existing nodes in the graph. Args: graph: The graph to get fake_mode from data: The tensor data to create fake tensor for Returns: A fake tensor with the appropriate fake_mode Raises:
(graph: torch.fx.Graph, data: torch.Tensor)
| 29 | |
| 30 | |
| 31 | def _get_fake_tensor_mode(graph: torch.fx.Graph, data: torch.Tensor) -> torch.Tensor: |
| 32 | """ |
| 33 | Helper function to create a fake tensor using the fake_mode from existing nodes in the graph. |
| 34 | |
| 35 | Args: |
| 36 | graph: The graph to get fake_mode from |
| 37 | data: The tensor data to create fake tensor for |
| 38 | |
| 39 | Returns: |
| 40 | A fake tensor with the appropriate fake_mode |
| 41 | |
| 42 | Raises: |
| 43 | RuntimeError: If the graph has no nodes to extract fake_mode from |
| 44 | """ |
| 45 | nodes = list(graph.nodes) |
| 46 | if not nodes: |
| 47 | raise RuntimeError( |
| 48 | "Cannot create fake tensor: graph has no nodes to extract fake_mode from" |
| 49 | ) |
| 50 | |
| 51 | example_node = nodes[0] |
| 52 | if isinstance( |
| 53 | example_node.meta["val"], (tuple, torch.fx.immutable_collections.immutable_list) |
| 54 | ): |
| 55 | example_fake_tensor = example_node.meta["val"][0] |
| 56 | else: |
| 57 | example_fake_tensor = example_node.meta["val"] |
| 58 | |
| 59 | return FakeTensorConverter().from_real_tensor(example_fake_tensor.fake_mode, t=data) |
| 60 | |
| 61 | |
| 62 | def is_get_attr_node(node: torch.fx.Node) -> bool: |
no outgoing calls
no test coverage detected