Skip to content

Fix/jax analyses do not parse file header#218

Merged
jujaykka merged 5 commits intomainfrom
fix/jax_analyses_do_not_parse_file_header
Jul 2, 2025
Merged

Fix/jax analyses do not parse file header#218
jujaykka merged 5 commits intomainfrom
fix/jax_analyses_do_not_parse_file_header

Conversation

@jujaykka
Copy link
Copy Markdown
Contributor

@jujaykka jujaykka commented Jul 1, 2025

When profiling a jit-compiled function like the train_step in Maxtext, the parsed xplane.pb will have a header line starting with "HloModule ". This line has no information that we currently want to parse but at the same time, it causes an exception to be raised. This commit completely ignores the header line.

olehtika and others added 5 commits June 25, 2025 13:28
- When profiling a jit-compiled function like the train_step in
  Maxtext, the parsed xplane.pb will have a header line starting with
  "HloModule ". This line has no information that we currently want to
  parse but at the same time, it causes an exception to be raised. This
  commit completely ignores the header line.
@jujaykka jujaykka requested review from gabeweisz and jasainio July 1, 2025 10:37
Copy link
Copy Markdown

@jasainio jasainio left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me. Rebase to main to make this cleaner once #210 is merged.

Copy link
Copy Markdown
Collaborator

@gabeweisz gabeweisz left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems to stack on top of 210. Looks fine to me - line 506 of jax_analyses.py is the only real change

@jujaykka jujaykka merged commit d8df0cc into main Jul 2, 2025
@jujaykka jujaykka deleted the fix/jax_analyses_do_not_parse_file_header branch July 2, 2025 11:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants