Return the shape of the tensor corresponding to node.
(
graph_module: torch.fx.GraphModule, node: torch.fx.Node
)
| 39 | |
| 40 | |
| 41 | def get_shape( |
| 42 | graph_module: torch.fx.GraphModule, node: torch.fx.Node |
| 43 | ) -> Union[torch.Size, None]: |
| 44 | """Return the shape of the tensor corresponding to node.""" |
| 45 | try: |
| 46 | if isinstance(node, (float, int, bool)): |
| 47 | return torch.Size([1]) |
| 48 | fake_tensor = node.meta.get("val") |
| 49 | if fake_tensor is not None: |
| 50 | return fake_tensor.shape |
| 51 | if node.op == "get_attr": |
| 52 | attr_node = getattr(graph_module, node.target) |
| 53 | return attr_node.shape |
| 54 | return None |
| 55 | except RuntimeError: |
| 56 | return None |
| 57 | |
| 58 | |
| 59 | def get_transposed_dims( |
no test coverage detected