Returns the number of tensors associated a given node
(node: torch.fx.Node)
| 291 | |
| 292 | |
| 293 | def num_tensors_in_node(node: torch.fx.Node) -> int: |
| 294 | """ |
| 295 | Returns the number of tensors associated a given node |
| 296 | """ |
| 297 | if "val" not in node.meta: |
| 298 | return 0 |
| 299 | |
| 300 | if isinstance(node.meta["val"], FakeTensor): |
| 301 | return 1 |
| 302 | |
| 303 | if isinstance(node.meta["val"], list) or isinstance(node.meta["val"], tuple): |
| 304 | if all(isinstance(x, FakeTensor) for x in node.meta["val"]): |
| 305 | return len(node.meta["val"]) |
| 306 | |
| 307 | return 0 |
| 308 | |
| 309 | |
| 310 | def get_vk_datatype(torch_dtype: torch.dtype) -> VkDataType: |
no outgoing calls
no test coverage detected