Checks whether the tensors produced by the given node can fit within the device's GPU buffer limit, which represents the maximum number of elements that can be stored in a GPU buffer.
(node: torch.fx.Node, buffer_limit: int)
| 700 | |
| 701 | |
| 702 | def within_buffer_limit(node: torch.fx.Node, buffer_limit: int) -> int: |
| 703 | """ |
| 704 | Checks whether the tensors produced by the given node can fit within the device's |
| 705 | GPU buffer limit, which represents the maximum number of elements that can be stored |
| 706 | in a GPU buffer. |
| 707 | """ |
| 708 | assert is_tensor_node(node) |
| 709 | |
| 710 | if isinstance(node.meta["val"], FakeTensor): |
| 711 | return node.meta["val"].numel() < buffer_limit |
| 712 | elif isinstance(node.meta["val"], list) or isinstance(node.meta["val"], tuple): |
| 713 | return all(x.numel() < buffer_limit for x in node.meta["val"]) |
| 714 | else: |
| 715 | raise RuntimeError(f"Cannot get numel for val of type {type(node.meta['val'])}") |
| 716 | |
| 717 | |
| 718 | def tensor_node_is_high_dim(node: torch.fx.Node) -> bool: |
nothing calls this directly
no test coverage detected