Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Compute source maps when pretty-printing jaxprs. #21081

Merged
merged 1 commit into from
May 7, 2024

Conversation

hawkinsp
Copy link
Collaborator

@hawkinsp hawkinsp commented May 6, 2024

This change is in preparation for adding support for emitting https://tc39.es/source-map/ source map information for jaxprs, so the relationship between a jaxpr and its Python code can be visualized using tooling for that purpose.

This change adds a new source_map() pretty printer document, which causes the pretty-printer to populate a source_map side output during pretty printing. The change also teaches the core jaxpr pretty printer to populate source map information on each equation.

@gspschmid

@hawkinsp hawkinsp requested a review from superbobry May 6, 2024 14:03
@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels May 6, 2024
@hawkinsp hawkinsp force-pushed the sourcemap branch 2 times, most recently from 3bfd69e to 4c4d1dd Compare May 6, 2024 18:33

Verified

This commit was signed with the committer’s verified signature.
Byron Sebastian Thiel
This change is in preparation for adding support for emitting https://tc39.es/source-map/ source map information for jaxprs, so the relationship between a jaxpr and its Python code can be visualized using tooling for that purpose.

This change adds a new `source_map()` pretty printer document, which causes the pretty-printer to populate a source_map side output during pretty printing.
The change also teaches the core jaxpr pretty printer to populate source map information on each equation.
@copybara-service copybara-service bot merged commit cb0c498 into jax-ml:main May 7, 2024
13 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

4 participants