Get a count of each MLX op node type in a serialized .pte file. Args: pte_path: Path to the .pte file Returns: Dictionary mapping op name (e.g. "SdpaNode", "SliceUpdateNode") to count.
(pte_path: Union[str, Path])
| 399 | |
| 400 | |
| 401 | def get_mlx_node_counts(pte_path: Union[str, Path]) -> Dict[str, int]: |
| 402 | """ |
| 403 | Get a count of each MLX op node type in a serialized .pte file. |
| 404 | |
| 405 | Args: |
| 406 | pte_path: Path to the .pte file |
| 407 | |
| 408 | Returns: |
| 409 | Dictionary mapping op name (e.g. "SdpaNode", "SliceUpdateNode") to count. |
| 410 | """ |
| 411 | data = inspect_pte_file(pte_path) |
| 412 | graph = data.get("graph", {}) |
| 413 | counts: Dict[str, int] = {} |
| 414 | for chain_info in graph.get("instruction_chains", []): |
| 415 | for instr in chain_info.get("instructions", []): |
| 416 | op_name = instr.get("op_name") |
| 417 | if op_name: |
| 418 | counts[op_name] = counts.get(op_name, 0) + 1 |
| 419 | return counts |
| 420 | |
| 421 | |
| 422 | def compare_outputs( |
no test coverage detected