Save kernel annotations to a pickle file.
(output_dir)
| 278 | # can find it. |
| 279 | |
| 280 | def save_annotations(output_dir): |
| 281 | """Save kernel annotations to a pickle file.""" |
| 282 | output_dir = Path(output_dir) |
| 283 | output_dir.mkdir(exist_ok=True, parents=True) |
| 284 | annotations_path = output_dir / "kernel_annotations_rank0_fwd_bwd.pkl" |
| 285 | |
| 286 | annotations = dict(get_kernel_annotations()) |
| 287 | with open(annotations_path, "wb") as f: |
| 288 | pickle.dump(annotations, f) |
| 289 | |
| 290 | print(f"Saved {len(annotations)} annotations to {annotations_path}") |
| 291 | return annotations_path |
| 292 | |
| 293 | ############################################################################### |
| 294 | # Post-Processing: Merging Annotations into Traces |
no outgoing calls
no test coverage detected